Add test for asserting measured cpu

databuild/job/BUILD.bazel

@@ -22,5 +22,6 @@ rust_test(
         "@crates//:serde_json",
         "@crates//:uuid",
         "@crates//:sysinfo",
+        "@crates//:tempfile",
     ],
 )

databuild/job/main.rs

@@ -2,29 +2,56 @@ use std::env;
 use std::io::{self, Read, Write};
 use std::process::{Command, Stdio};
 use std::sync::{Arc, Mutex};
-use std::time::{SystemTime, UNIX_EPOCH, Duration};
 use std::thread;
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
 // All serialization handled by protobuf serde derives
 use serde_json;
+use sysinfo::{Pid, ProcessRefreshKind, System};
 use uuid::Uuid;
-use sysinfo::{System, ProcessRefreshKind, Pid};
 
 // Import protobuf types from databuild
 use databuild::{
-    PartitionRef, PartitionManifest, Task, JobLabel, JobConfig,
+    job_log_entry, log_message, JobConfig, JobLabel, JobLogEntry, LogMessage, PartitionManifest,
-    JobLogEntry, LogMessage, WrapperJobEvent, job_log_entry, log_message
+    PartitionRef, Task, WrapperJobEvent,
 };
 
 // All types now come from protobuf - no custom structs needed
 
-struct JobWrapper {
+fn get_timestamp() -> String {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_secs()
+        .to_string()
+}
+
+trait LogSink {
+    fn emit(&mut self, entry: JobLogEntry);
+}
+
+struct StdoutSink;
+
+impl LogSink for StdoutSink {
+    fn emit(&mut self, entry: JobLogEntry) {
+        println!("{}", serde_json::to_string(&entry).unwrap());
+    }
+}
+
+struct JobWrapper<S: LogSink> {
     job_id: String,
     sequence_number: u64,
     start_time: i64,
+    sink: S,
 }
 
-impl JobWrapper {
+impl JobWrapper<StdoutSink> {
     fn new() -> Self {
+        Self::new_with_sink(StdoutSink)
+    }
+}
+
+impl<S: LogSink> JobWrapper<S> {
+    fn new_with_sink(sink: S) -> Self {
         Self {
             job_id: Uuid::new_v4().to_string(),
             sequence_number: 0,
@@ -32,17 +59,10 @@ impl JobWrapper {
                 .duration_since(UNIX_EPOCH)
                 .unwrap()
                 .as_secs() as i64,
+            sink,
         }
     }
 
-    fn get_timestamp() -> String {
-        SystemTime::now()
-            .duration_since(UNIX_EPOCH)
-            .unwrap()
-            .as_secs()
-            .to_string()
-    }
-
     fn next_sequence(&mut self) -> u64 {
         self.sequence_number += 1;
         self.sequence_number
@@ -50,14 +70,14 @@ impl JobWrapper {
 
     fn emit_log(&mut self, partition_ref: &str, content: job_log_entry::Content) {
         let entry = JobLogEntry {
-            timestamp: Self::get_timestamp(),
+            timestamp: get_timestamp(),
             job_id: self.job_id.clone(),
             partition_ref: partition_ref.to_string(),
             sequence_number: self.next_sequence(),
             content: Some(content),
         };
 
-        println!("{}", serde_json::to_string(&entry).unwrap());
+        self.sink.emit(entry);
     }
 
     fn config_mode(&mut self, outputs: Vec<String>) -> Result<(), Box<dyn std::error::Error>> {
@@ -65,17 +85,23 @@ impl JobWrapper {
         let partition_ref = outputs.first().unwrap_or(&"unknown".to_string()).clone();
 
         // Following the state diagram: wrapper_validate_config -> emit_config_validate_success
-        self.emit_log(&partition_ref, job_log_entry::Content::JobEvent(WrapperJobEvent {
+        self.emit_log(
+            &partition_ref,
+            job_log_entry::Content::JobEvent(WrapperJobEvent {
                 event_type: "config_validate_success".to_string(),
                 metadata: std::collections::HashMap::new(),
                 job_status: None,
                 exit_code: None,
-        }));
+            }),
+        );
 
         // For Phase 0, we still need to produce the expected JSON config format
         // so the current graph system can parse it. Later phases will change this.
         let config = JobConfig {
-            outputs: outputs.iter().map(|s| PartitionRef { r#str: s.clone() }).collect(),
+            outputs: outputs
+                .iter()
+                .map(|s| PartitionRef { r#str: s.clone() })
+                .collect(),
             inputs: vec![],
             args: outputs.clone(),
             env: {
@@ -104,26 +130,42 @@ impl JobWrapper {
         io::stdin().read_to_string(&mut buffer)?;
 
         let config: JobConfig = serde_json::from_str(&buffer)?;
-        let partition_ref = config.outputs.first()
+        self.exec_mode_with_config(job_binary, config)
+    }
+
+    fn exec_mode_with_config(
+        &mut self,
+        job_binary: &str,
+        config: JobConfig,
+    ) -> Result<(), Box<dyn std::error::Error>> {
+        let partition_ref = config
+            .outputs
+            .first()
             .map(|p| p.str.clone())
             .unwrap_or_else(|| "unknown".to_string());
 
         // Following the state diagram:
         // 1. wrapper_validate_config -> emit_config_validate_success
-        self.emit_log(&partition_ref, job_log_entry::Content::JobEvent(WrapperJobEvent {
+        self.emit_log(
+            &partition_ref,
+            job_log_entry::Content::JobEvent(WrapperJobEvent {
                 event_type: "config_validate_success".to_string(),
                 job_status: None,
                 exit_code: None,
                 metadata: std::collections::HashMap::new(),
-        }));
+            }),
+        );
 
         // 2. wrapper_launch_task -> emit_task_launch_success
-        self.emit_log(&partition_ref, job_log_entry::Content::JobEvent(WrapperJobEvent {
+        self.emit_log(
+            &partition_ref,
+            job_log_entry::Content::JobEvent(WrapperJobEvent {
                 event_type: "task_launch_success".to_string(),
                 job_status: None,
                 exit_code: None,
                 metadata: std::collections::HashMap::new(),
-        }));
+            }),
+        );
 
         // Execute the original job binary with the exec subcommand
         let mut cmd = Command::new(job_binary);
@@ -148,7 +190,7 @@ impl JobWrapper {
 
         // Send the config to the job
         if let Some(stdin) = child.stdin.as_mut() {
-            stdin.write_all(buffer.as_bytes())?;
+            stdin.write_all(serde_json::to_string(&config).unwrap().as_bytes())?;
         }
 
         // Start heartbeat thread
@@ -175,7 +217,10 @@ impl JobWrapper {
                     // Create heartbeat event with metrics
                     let mut metadata = std::collections::HashMap::new();
                     metadata.insert("memory_usage_mb".to_string(), format!("{:.3}", memory_mb));
-                    metadata.insert("cpu_usage_percent".to_string(), format!("{:.3}", cpu_percent));
+                    metadata.insert(
+                        "cpu_usage_percent".to_string(),
+                        format!("{:.3}", cpu_percent),
+                    );
 
                     // Get next sequence number for heartbeat
                     let seq = {
@@ -185,7 +230,7 @@ impl JobWrapper {
                     };
 
                     let heartbeat_event = JobLogEntry {
-                        timestamp: JobWrapper::get_timestamp(),
+                        timestamp: get_timestamp(),
                         job_id: heartbeat_job_id.clone(),
                         partition_ref: heartbeat_partition_ref.clone(),
                         sequence_number: seq,
@@ -230,14 +275,19 @@ impl JobWrapper {
                     }
 
                     // Calculate final metrics
-                    let job_duration = job_start_time.elapsed()
+                    let job_duration = job_start_time.elapsed().map_err(|e| {
-                        .map_err(|e| io::Error::new(io::ErrorKind::Other, format!("Time calculation error: {}", e)))?;
+                        io::Error::new(
+                            io::ErrorKind::Other,
+                            format!("Time calculation error: {}", e),
+                        )
+                    })?;
 
                     // Calculate CPU time: average CPU percentage * wall-clock time
                     let total_cpu_ms = if cpu_samples.is_empty() {
                         0.0
                     } else {
-                        let avg_cpu_percent = cpu_samples.iter().sum::<f32>() as f64 / cpu_samples.len() as f64;
+                        let avg_cpu_percent =
+                            cpu_samples.iter().sum::<f32>() as f64 / cpu_samples.len() as f64;
                         (avg_cpu_percent / 100.0) * job_duration.as_millis() as f64
                     };
 
@@ -283,44 +333,62 @@ impl JobWrapper {
         // Capture and forward job stdout/stderr as log messages
         if !output.stdout.is_empty() {
             let stdout_str = String::from_utf8_lossy(&output.stdout);
-            self.emit_log(&partition_ref, job_log_entry::Content::Log(LogMessage {
+            self.emit_log(
+                &partition_ref,
+                job_log_entry::Content::Log(LogMessage {
                     level: log_message::LogLevel::Info as i32,
                     message: stdout_str.to_string(),
                     fields: std::collections::HashMap::new(),
-            }));
+                }),
+            );
         }
 
         if !output.stderr.is_empty() {
             let stderr_str = String::from_utf8_lossy(&output.stderr);
-            self.emit_log(&partition_ref, job_log_entry::Content::Log(LogMessage {
+            self.emit_log(
+                &partition_ref,
+                job_log_entry::Content::Log(LogMessage {
                     level: log_message::LogLevel::Error as i32,
                     message: stderr_str.to_string(),
                     fields: std::collections::HashMap::new(),
-            }));
+                }),
+            );
         }
 
         // Emit job summary with resource metrics
         let mut summary_metadata = std::collections::HashMap::new();
-        summary_metadata.insert("runtime_ms".to_string(), format!("{:.3}", job_duration.as_millis() as f64));
+        summary_metadata.insert(
-        summary_metadata.insert("peak_memory_mb".to_string(), format!("{:.3}", peak_memory_mb));
+            "runtime_ms".to_string(),
+            format!("{:.3}", job_duration.as_millis() as f64),
+        );
+        summary_metadata.insert(
+            "peak_memory_mb".to_string(),
+            format!("{:.3}", peak_memory_mb),
+        );
         summary_metadata.insert("total_cpu_ms".to_string(), format!("{:.3}", total_cpu_ms));
         summary_metadata.insert("exit_code".to_string(), exit_code.to_string());
 
-        self.emit_log(&partition_ref, job_log_entry::Content::JobEvent(WrapperJobEvent {
+        self.emit_log(
+            &partition_ref,
+            job_log_entry::Content::JobEvent(WrapperJobEvent {
                 event_type: "job_summary".to_string(),
                 job_status: None,
                 exit_code: Some(exit_code),
                 metadata: summary_metadata,
-        }));
+            }),
+        );
 
         if success {
             // Following the state diagram: wrapper_monitor_task -> zero exit -> emit_task_success
-            self.emit_log(&partition_ref, job_log_entry::Content::JobEvent(WrapperJobEvent {
+            self.emit_log(
+                &partition_ref,
+                job_log_entry::Content::JobEvent(WrapperJobEvent {
                     event_type: "task_success".to_string(),
                     job_status: Some("JOB_COMPLETED".to_string()),
                     exit_code: Some(exit_code),
                     metadata: std::collections::HashMap::new(),
-            }));
+                }),
+            );
 
             // Then emit_partition_manifest -> success
             let end_time = SystemTime::now()
@@ -328,39 +396,52 @@ impl JobWrapper {
                 .unwrap()
                 .as_secs() as i64;
 
-            self.emit_log(&partition_ref, job_log_entry::Content::Manifest(PartitionManifest {
+            self.emit_log(
+                &partition_ref,
+                job_log_entry::Content::Manifest(PartitionManifest {
                     outputs: config.outputs.clone(),
                     inputs: vec![], // Phase 0: no input manifests yet
                     start_time: self.start_time,
                     end_time,
                     task: Some(Task {
                         job: Some(JobLabel {
-                        label: env::var("DATABUILD_JOB_LABEL").unwrap_or_else(|_| "unknown".to_string()),
+                            label: env::var("DATABUILD_JOB_LABEL")
+                                .unwrap_or_else(|_| "unknown".to_string()),
                         }),
                         config: Some(config.clone()),
                     }),
                     metadata: std::collections::HashMap::new(), // Phase 0: no metadata yet
-            }));
+                }),
+            );
         } else {
             // Following the state diagram: wrapper_monitor_task -> non-zero exit -> emit_task_failed
-            self.emit_log(&partition_ref, job_log_entry::Content::JobEvent(WrapperJobEvent {
+            self.emit_log(
+                &partition_ref,
+                job_log_entry::Content::JobEvent(WrapperJobEvent {
                     event_type: "task_failed".to_string(),
                     job_status: Some("JOB_FAILED".to_string()),
                     exit_code: Some(exit_code),
                     metadata: std::collections::HashMap::new(),
-            }));
+                }),
+            );
 
             // Then emit_job_exec_fail -> fail (don't emit partition manifest on failure)
-            self.emit_log(&partition_ref, job_log_entry::Content::JobEvent(WrapperJobEvent {
+            self.emit_log(
+                &partition_ref,
+                job_log_entry::Content::JobEvent(WrapperJobEvent {
                     event_type: "job_exec_fail".to_string(),
                     job_status: Some("JOB_FAILED".to_string()),
                     exit_code: Some(exit_code),
                     metadata: {
                         let mut meta = std::collections::HashMap::new();
-                    meta.insert("error".to_string(), format!("Job failed with exit code {}", exit_code));
+                        meta.insert(
+                            "error".to_string(),
+                            format!("Job failed with exit code {}", exit_code),
+                        );
                         meta
                     },
-            }));
+                }),
+            );
         }
 
         // Forward the original job's output to stdout for compatibility
@@ -394,8 +475,8 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         "exec" => {
             // For exec mode, we need to know which original job binary to call
             // For Phase 0, we'll derive this from environment or make it configurable
-            let job_binary = env::var("DATABUILD_JOB_BINARY")
+            let job_binary =
-                .unwrap_or_else(|_| "python3".to_string()); // Default fallback
+                env::var("DATABUILD_JOB_BINARY").unwrap_or_else(|_| "python3".to_string()); // Default fallback
 
             wrapper.exec_mode(&job_binary)?;
         }
@@ -412,10 +493,42 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 mod tests {
     use super::*;
 
+    // Test infrastructure
+    struct TestSink {
+        entries: Vec<JobLogEntry>,
+    }
+
+    impl TestSink {
+        fn new() -> Self {
+            Self {
+                entries: Vec::new(),
+            }
+        }
+
+        fn find_event(&self, event_type: &str) -> Option<&JobLogEntry> {
+            self.entries.iter().find(|entry| {
+                if let Some(job_log_entry::Content::JobEvent(event)) = &entry.content {
+                    event.event_type == event_type
+                } else {
+                    false
+                }
+            })
+        }
+    }
+
+    impl LogSink for TestSink {
+        fn emit(&mut self, entry: JobLogEntry) {
+            self.entries.push(entry);
+        }
+    }
+
     // Helper functions for testing
     fn generate_test_config(outputs: &[String]) -> JobConfig {
         JobConfig {
-            outputs: outputs.iter().map(|s| PartitionRef { r#str: s.clone() }).collect(),
+            outputs: outputs
+                .iter()
+                .map(|s| PartitionRef { r#str: s.clone() })
+                .collect(),
             inputs: vec![],
             args: outputs.to_vec(),
             env: {
@@ -466,7 +579,10 @@ mod tests {
         assert_eq!(config.outputs.len(), 1);
         assert_eq!(config.outputs[0].r#str, "test/partition");
         assert_eq!(config.args, outputs);
-        assert_eq!(config.env.get("PARTITION_REF"), Some(&"test/partition".to_string()));
+        assert_eq!(
+            config.env.get("PARTITION_REF"),
+            Some(&"test/partition".to_string())
+        );
     }
 
     #[test]
@@ -481,7 +597,10 @@ mod tests {
         assert_eq!(config.outputs[0].r#str, "reviews/date=2025-01-01");
         assert_eq!(config.outputs[1].r#str, "reviews/date=2025-01-02");
         // First output is used as PARTITION_REF
-        assert_eq!(config.env.get("PARTITION_REF"), Some(&"reviews/date=2025-01-01".to_string()));
+        assert_eq!(
+            config.env.get("PARTITION_REF"),
+            Some(&"reviews/date=2025-01-01".to_string())
+        );
     }
 
     #[test]
@@ -551,9 +670,9 @@ mod tests {
 
     #[test]
     fn test_timestamp_generation() {
-        let ts1 = JobWrapper::get_timestamp();
+        let ts1 = get_timestamp();
         std::thread::sleep(std::time::Duration::from_millis(10));
-        let ts2 = JobWrapper::get_timestamp();
+        let ts2 = get_timestamp();
 
         // Timestamps should be parseable as integers
         let t1: u64 = ts1.parse().expect("Should be valid timestamp");
@@ -573,75 +692,117 @@ mod tests {
 
     #[test]
     fn test_cpu_metrics_are_captured() {
-        use std::process::{Command, Stdio};
         use std::io::Write;
+        use std::process::{Command, Stdio};
+        use tempfile::NamedTempFile;
 
-        // Set faster sampling interval for test
+        // Create a CPU-intensive test script
-        env::set_var("DATABUILD_METRICS_INTERVAL_MS", "10");
+        let mut temp_file = NamedTempFile::new().expect("Failed to create temp file");
-        
+        let script_content = r#"#!/usr/bin/env python3
-        // Create a CPU-intensive Python script
-        let cpu_script = r#"
 import sys
+import json
+
 if len(sys.argv) > 1 and sys.argv[1] == "config":
-    print('{"configs":[{"outputs":[{"str":"test/cpu"}],"inputs":[],"args":[],"env":{}}]}')
+    config = {
+        "outputs": [{"str": "test/cpu"}],
+        "inputs": [],
+        "args": [],
+        "env": {"PARTITION_REF": "test/cpu"}
+    }
+    print(json.dumps({"configs": [config]}))
 elif len(sys.argv) > 1 and sys.argv[1] == "exec":
     # CPU-intensive work
-    total = sum(range(10_000_000))
+    total = sum(range(50_000_000))  # Smaller for faster test
     print(f"Sum: {total}")
 "#;
 
-        let script_path = "/tmp/databuild_test_cpu.py";
+        temp_file
-        std::fs::write(script_path, cpu_script).expect("Failed to write test script");
+            .write_all(script_content.as_bytes())
+            .expect("Failed to write script");
+        let script_path = temp_file.path().to_str().unwrap();
 
-        // Run wrapper with the CPU-intensive job
+        // Make script executable
-        env::set_var("DATABUILD_JOB_BINARY", "python3");
+        std::fs::set_permissions(
+            script_path,
+            std::os::unix::fs::PermissionsExt::from_mode(0o755),
+        )
+        .expect("Failed to set permissions");
 
-        let output = Command::new(env!("CARGO_BIN_EXE_job_wrapper"))
+        // Set up environment for fast sampling and the test script
-            .arg("exec")
+        env::set_var("DATABUILD_METRICS_INTERVAL_MS", "10");
-            .stdin(Stdio::piped())
+        env::set_var("DATABUILD_JOB_BINARY", script_path);
-            .stdout(Stdio::piped())
+
-            .stderr(Stdio::piped())
+        // Create test sink and wrapper
-            .spawn()
+        let sink = TestSink::new();
-            .and_then(|mut child| {
+        let mut wrapper = JobWrapper::new_with_sink(sink);
-                // Send config to stdin
+
+        // Create a JobConfig for the test
         let config = JobConfig {
-                    outputs: vec![PartitionRef { r#str: "test/cpu".to_string() }],
+            outputs: vec![PartitionRef {
+                r#str: "test/cpu".to_string(),
+            }],
             inputs: vec![],
-                    args: vec![script_path.to_string()],
+            args: vec![],
-                    env: std::collections::HashMap::new(),
+            env: {
+                let mut env_map = std::collections::HashMap::new();
+                env_map.insert("PARTITION_REF".to_string(), "test/cpu".to_string());
+                env_map
+            },
         };
-                let config_json = serde_json::to_string(&config).unwrap();
 
-                if let Some(stdin) = child.stdin.as_mut() {
+        // We need to simulate stdin for exec_mode - let's create a test-specific exec method
-                    stdin.write_all(config_json.as_bytes()).unwrap();
+        // that takes the config directly rather than reading from stdin
-                }
+        let result = wrapper.exec_mode_with_config(script_path, config);
 
-                child.wait_with_output()
+        // Clean up environment
-            })
-            .expect("Failed to run wrapper");
-        
-        // Parse output to find job_summary
-        let stdout = String::from_utf8_lossy(&output.stdout);
-        let mut cpu_ms = None;
-        
-        for line in stdout.lines() {
-            if let Ok(json) = serde_json::from_str::<serde_json::Value>(line) {
-                if json["content"]["JobEvent"]["event_type"] == "job_summary" {
-                    if let Some(cpu_str) = json["content"]["JobEvent"]["metadata"]["total_cpu_ms"].as_str() {
-                        cpu_ms = Some(cpu_str.parse::<f64>().unwrap());
-                    }
-                }
-            }
-        }
-        
-        // Clean up
-        std::fs::remove_file(script_path).ok();
         env::remove_var("DATABUILD_METRICS_INTERVAL_MS");
+        env::remove_var("DATABUILD_JOB_BINARY");
 
-        // Assert we captured non-zero CPU time
+        // Check that exec_mode succeeded
-        assert!(cpu_ms.is_some(), "Should have found job_summary event");
+        if let Err(e) = &result {
-        let cpu_time = cpu_ms.unwrap();
+            println!("exec_mode failed with error: {}", e);
-        assert!(cpu_time > 0.0, 
+        }
-            "Expected non-zero CPU time for CPU-intensive workload, but got {:.3}ms", cpu_time);
+        assert!(result.is_ok(), "exec_mode should succeed: {:?}", result);
+
+        // Find the job_summary event
+        let summary_event = wrapper
+            .sink
+            .find_event("job_summary")
+            .expect("Should have job_summary event");
+
+        if let Some(job_log_entry::Content::JobEvent(event)) = &summary_event.content {
+            // Verify we have CPU metrics
+            let cpu_ms_str = event
+                .metadata
+                .get("total_cpu_ms")
+                .expect("Should have total_cpu_ms metric");
+            let cpu_ms: f64 = cpu_ms_str
+                .parse()
+                .expect("CPU metric should be valid float");
+
+            // For CPU-intensive work, we should get non-zero CPU time
+            assert!(
+                cpu_ms > 0.0,
+                "Expected non-zero CPU time for CPU-intensive workload, but got {:.3}ms",
+                cpu_ms
+            );
+
+            // Also verify runtime is reasonable
+            let runtime_ms_str = event
+                .metadata
+                .get("runtime_ms")
+                .expect("Should have runtime_ms metric");
+            let runtime_ms: f64 = runtime_ms_str
+                .parse()
+                .expect("Runtime metric should be valid float");
+            assert!(runtime_ms > 0.0, "Should have non-zero runtime");
+
+            println!(
+                "CPU test results: {:.3}ms CPU time over {:.3}ms runtime",
+                cpu_ms, runtime_ms
+            );
+        } else {
+            panic!("job_summary event should contain JobEvent");
+        }
     }
 }