From 31db6a00cbb0d1992fdedb89303456d537984597 Mon Sep 17 00:00:00 2001 From: Stuart Axelbrooke Date: Tue, 25 Nov 2025 14:03:38 +0800 Subject: [PATCH] add tests --- databuild/build_state.rs | 378 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 378 insertions(+) diff --git a/databuild/build_state.rs b/databuild/build_state.rs index d839b26..3e5d7ae 100644 --- a/databuild/build_state.rs +++ b/databuild/build_state.rs @@ -1957,6 +1957,384 @@ mod tests { "Alpha want should be successful" ); } + + /// Test that multiple concurrent wants for the same partition all transition correctly. + /// This was the original bug that motivated the UUID refactor. + #[test] + fn test_concurrent_wants_same_partition() { + use crate::data_build_event::Event; + use crate::{ + JobRunBufferEventV1, JobRunHeartbeatEventV1, PartitionRef, + WantAttributedPartitions, WantCreateEventV1, + }; + + let mut state = BuildState::default(); + + // 1. Create Want 1 for data/beta + let want_1_id = "want-1".to_string(); + let mut create_want_1 = WantCreateEventV1::default(); + create_want_1.want_id = want_1_id.clone(); + create_want_1.partitions = vec![PartitionRef { + r#ref: "data/beta".to_string(), + }]; + state.handle_event(&Event::WantCreateV1(create_want_1)); + + // Want 1 should be Idle (no partition exists yet) + let want_1 = state.get_want(&want_1_id).unwrap(); + assert_eq!( + want_1.status, + Some(crate::WantStatusCode::WantIdle.into()), + "Want 1 should be Idle initially" + ); + + // 2. Create Want 2 for the same partition + let want_2_id = "want-2".to_string(); + let mut create_want_2 = WantCreateEventV1::default(); + create_want_2.want_id = want_2_id.clone(); + create_want_2.partitions = vec![PartitionRef { + r#ref: "data/beta".to_string(), + }]; + state.handle_event(&Event::WantCreateV1(create_want_2)); + + // Want 2 should also be Idle + let want_2 = state.get_want(&want_2_id).unwrap(); + assert_eq!( + want_2.status, + Some(crate::WantStatusCode::WantIdle.into()), + "Want 2 should be Idle initially" + ); + + // Verify inverted index has both wants + let wants_for_beta = state.get_wants_for_partition("data/beta"); + assert!( + wants_for_beta.contains(&want_1_id), + "wants_for_partition should contain want-1" + ); + assert!( + wants_for_beta.contains(&want_2_id), + "wants_for_partition should contain want-2" + ); + + // 3. Job buffers for data/beta - both wants should transition to Building + let job_run_id = "job-1".to_string(); + let mut buffer_event = JobRunBufferEventV1::default(); + buffer_event.job_run_id = job_run_id.clone(); + buffer_event.job_label = "//job_beta".to_string(); + buffer_event.want_attributed_partitions = vec![ + WantAttributedPartitions { + want_id: want_1_id.clone(), + partitions: vec![PartitionRef { + r#ref: "data/beta".to_string(), + }], + }, + WantAttributedPartitions { + want_id: want_2_id.clone(), + partitions: vec![PartitionRef { + r#ref: "data/beta".to_string(), + }], + }, + ]; + buffer_event.building_partitions = vec![PartitionRef { + r#ref: "data/beta".to_string(), + }]; + state.handle_event(&Event::JobRunBufferV1(buffer_event)); + + // Start the job + let mut heartbeat = JobRunHeartbeatEventV1::default(); + heartbeat.job_run_id = job_run_id.clone(); + state.handle_event(&Event::JobRunHeartbeatV1(heartbeat)); + + // Both wants should now be Building + let want_1 = state.get_want(&want_1_id).unwrap(); + assert_eq!( + want_1.status, + Some(crate::WantStatusCode::WantBuilding.into()), + "Want 1 should be Building after job starts" + ); + + let want_2 = state.get_want(&want_2_id).unwrap(); + assert_eq!( + want_2.status, + Some(crate::WantStatusCode::WantBuilding.into()), + "Want 2 should be Building after job starts" + ); + + // Partition should exist and be Building + let partition = state.get_partition("data/beta").unwrap(); + assert_eq!( + partition.status, + Some(crate::PartitionStatusCode::PartitionBuilding.into()), + "Partition should be Building" + ); + } + } + + mod partition_lifecycle { + use crate::build_state::BuildState; + use crate::data_build_event::Event; + use crate::{ + JobRunBufferEventV1, JobRunFailureEventV1, JobRunHeartbeatEventV1, + JobRunMissingDepsEventV1, JobRunSuccessEventV1, MissingDeps, PartitionRef, + WantAttributedPartitions, WantCreateEventV1, + }; + + /// Test that upstream failure cascades to downstream partitions. + /// When an upstream partition fails, downstream partitions in UpstreamBuilding + /// should transition to UpstreamFailed. + #[test] + fn test_upstream_failure_cascades_to_downstream() { + let mut state = BuildState::default(); + + // 1. Create want for data/beta + let beta_want_id = "beta-want".to_string(); + let mut create_beta = WantCreateEventV1::default(); + create_beta.want_id = beta_want_id.clone(); + create_beta.partitions = vec![PartitionRef { + r#ref: "data/beta".to_string(), + }]; + state.handle_event(&Event::WantCreateV1(create_beta)); + + // 2. Job buffers for beta + let beta_job_id = "beta-job".to_string(); + let mut buffer_beta = JobRunBufferEventV1::default(); + buffer_beta.job_run_id = beta_job_id.clone(); + buffer_beta.job_label = "//job_beta".to_string(); + buffer_beta.want_attributed_partitions = vec![WantAttributedPartitions { + want_id: beta_want_id.clone(), + partitions: vec![PartitionRef { + r#ref: "data/beta".to_string(), + }], + }]; + buffer_beta.building_partitions = vec![PartitionRef { + r#ref: "data/beta".to_string(), + }]; + state.handle_event(&Event::JobRunBufferV1(buffer_beta)); + + // Start beta job + let mut heartbeat_beta = JobRunHeartbeatEventV1::default(); + heartbeat_beta.job_run_id = beta_job_id.clone(); + state.handle_event(&Event::JobRunHeartbeatV1(heartbeat_beta)); + + // 3. Beta job dep misses on data/alpha + // This returns derivative want events that we MUST process + let mut dep_miss = JobRunMissingDepsEventV1::default(); + dep_miss.job_run_id = beta_job_id.clone(); + dep_miss.missing_deps = vec![MissingDeps { + impacted: vec![PartitionRef { + r#ref: "data/beta".to_string(), + }], + missing: vec![PartitionRef { + r#ref: "data/alpha".to_string(), + }], + }]; + let derivative_events = state.handle_event(&Event::JobRunMissingDepsV1(dep_miss)); + + // Beta partition should be UpstreamBuilding + let beta_partition = state.get_partition("data/beta").unwrap(); + assert_eq!( + beta_partition.status.as_ref().unwrap().name, + "PartitionUpstreamBuilding", + "Beta partition should be UpstreamBuilding after dep miss" + ); + + // 4. Process the derivative want event for alpha + // This will create the alpha want AND transition beta want to UpstreamBuilding + assert_eq!( + derivative_events.len(), + 1, + "Should have one derivative want event" + ); + let alpha_want_id = match &derivative_events[0] { + Event::WantCreateV1(e) => e.want_id.clone(), + _ => panic!("Expected WantCreateV1 event"), + }; + state.handle_event(&derivative_events[0]); + + // Now beta want should be UpstreamBuilding (waiting for alpha want) + let beta_want_after_derivative = state.get_want(&beta_want_id).unwrap(); + assert_eq!( + beta_want_after_derivative.status, + Some(crate::WantStatusCode::WantUpstreamBuilding.into()), + "Beta want should be UpstreamBuilding after derivative want processed" + ); + + // 5. Job buffers for alpha + let alpha_job_id = "alpha-job".to_string(); + let mut buffer_alpha = JobRunBufferEventV1::default(); + buffer_alpha.job_run_id = alpha_job_id.clone(); + buffer_alpha.job_label = "//job_alpha".to_string(); + buffer_alpha.want_attributed_partitions = vec![WantAttributedPartitions { + want_id: alpha_want_id.clone(), + partitions: vec![PartitionRef { + r#ref: "data/alpha".to_string(), + }], + }]; + buffer_alpha.building_partitions = vec![PartitionRef { + r#ref: "data/alpha".to_string(), + }]; + state.handle_event(&Event::JobRunBufferV1(buffer_alpha)); + + // Start alpha job + let mut heartbeat_alpha = JobRunHeartbeatEventV1::default(); + heartbeat_alpha.job_run_id = alpha_job_id.clone(); + state.handle_event(&Event::JobRunHeartbeatV1(heartbeat_alpha)); + + // 6. Alpha job FAILS + let mut failure_alpha = JobRunFailureEventV1::default(); + failure_alpha.job_run_id = alpha_job_id.clone(); + failure_alpha.reason = "Test failure".to_string(); + state.handle_event(&Event::JobRunFailureV1(failure_alpha)); + + // Alpha partition should be Failed + let alpha_partition = state.get_partition("data/alpha").unwrap(); + assert_eq!( + alpha_partition.status, + Some(crate::PartitionStatusCode::PartitionFailed.into()), + "Alpha partition should be Failed" + ); + + // Beta partition should cascade to UpstreamFailed + let beta_partition = state.get_partition("data/beta").unwrap(); + assert_eq!( + beta_partition.status.as_ref().unwrap().name, + "PartitionUpstreamFailed", + "Beta partition should cascade to UpstreamFailed when alpha fails" + ); + + // Beta want should also be UpstreamFailed + let beta_want = state.get_want(&beta_want_id).unwrap(); + assert_eq!( + beta_want.status, + Some(crate::WantStatusCode::WantUpstreamFailed.into()), + "Beta want should be UpstreamFailed" + ); + } + + /// Test that partition retry creates a new UUID while preserving the old one. + #[test] + fn test_partition_retry_creates_new_uuid() { + let mut state = BuildState::default(); + + // 1. Create want for data/beta + let beta_want_id = "beta-want".to_string(); + let mut create_beta = WantCreateEventV1::default(); + create_beta.want_id = beta_want_id.clone(); + create_beta.partitions = vec![PartitionRef { + r#ref: "data/beta".to_string(), + }]; + state.handle_event(&Event::WantCreateV1(create_beta)); + + // 2. First job buffers for beta (creates uuid-1) + let beta_job_1_id = "beta-job-1".to_string(); + let mut buffer_beta_1 = JobRunBufferEventV1::default(); + buffer_beta_1.job_run_id = beta_job_1_id.clone(); + buffer_beta_1.job_label = "//job_beta".to_string(); + buffer_beta_1.want_attributed_partitions = vec![WantAttributedPartitions { + want_id: beta_want_id.clone(), + partitions: vec![PartitionRef { + r#ref: "data/beta".to_string(), + }], + }]; + buffer_beta_1.building_partitions = vec![PartitionRef { + r#ref: "data/beta".to_string(), + }]; + state.handle_event(&Event::JobRunBufferV1(buffer_beta_1)); + + // Get UUID-1 + let partition_after_first_buffer = state.get_partition("data/beta").unwrap(); + let uuid_1 = partition_after_first_buffer.uuid.clone(); + + // Start job + let mut heartbeat_1 = JobRunHeartbeatEventV1::default(); + heartbeat_1.job_run_id = beta_job_1_id.clone(); + state.handle_event(&Event::JobRunHeartbeatV1(heartbeat_1)); + + // 3. Job dep misses on data/alpha + let mut dep_miss = JobRunMissingDepsEventV1::default(); + dep_miss.job_run_id = beta_job_1_id.clone(); + dep_miss.missing_deps = vec![MissingDeps { + impacted: vec![PartitionRef { + r#ref: "data/beta".to_string(), + }], + missing: vec![PartitionRef { + r#ref: "data/alpha".to_string(), + }], + }]; + state.handle_event(&Event::JobRunMissingDepsV1(dep_miss)); + + // 4. Create and complete alpha (to satisfy the dep) + let alpha_want_id = "alpha-want".to_string(); + let mut create_alpha = WantCreateEventV1::default(); + create_alpha.want_id = alpha_want_id.clone(); + create_alpha.partitions = vec![PartitionRef { + r#ref: "data/alpha".to_string(), + }]; + state.handle_event(&Event::WantCreateV1(create_alpha)); + + let alpha_job_id = "alpha-job".to_string(); + let mut buffer_alpha = JobRunBufferEventV1::default(); + buffer_alpha.job_run_id = alpha_job_id.clone(); + buffer_alpha.job_label = "//job_alpha".to_string(); + buffer_alpha.want_attributed_partitions = vec![WantAttributedPartitions { + want_id: alpha_want_id.clone(), + partitions: vec![PartitionRef { + r#ref: "data/alpha".to_string(), + }], + }]; + buffer_alpha.building_partitions = vec![PartitionRef { + r#ref: "data/alpha".to_string(), + }]; + state.handle_event(&Event::JobRunBufferV1(buffer_alpha)); + + let mut heartbeat_alpha = JobRunHeartbeatEventV1::default(); + heartbeat_alpha.job_run_id = alpha_job_id.clone(); + state.handle_event(&Event::JobRunHeartbeatV1(heartbeat_alpha)); + + let mut success_alpha = JobRunSuccessEventV1::default(); + success_alpha.job_run_id = alpha_job_id.clone(); + state.handle_event(&Event::JobRunSuccessV1(success_alpha)); + + // Beta partition should now be UpForRetry + let beta_partition = state.get_partition("data/beta").unwrap(); + assert_eq!( + beta_partition.status.as_ref().unwrap().name, + "PartitionUpForRetry", + "Beta partition should be UpForRetry after alpha succeeds" + ); + + // 5. Second job buffers for beta retry (creates uuid-2) + let beta_job_2_id = "beta-job-2".to_string(); + let mut buffer_beta_2 = JobRunBufferEventV1::default(); + buffer_beta_2.job_run_id = beta_job_2_id.clone(); + buffer_beta_2.job_label = "//job_beta".to_string(); + buffer_beta_2.want_attributed_partitions = vec![WantAttributedPartitions { + want_id: beta_want_id.clone(), + partitions: vec![PartitionRef { + r#ref: "data/beta".to_string(), + }], + }]; + buffer_beta_2.building_partitions = vec![PartitionRef { + r#ref: "data/beta".to_string(), + }]; + state.handle_event(&Event::JobRunBufferV1(buffer_beta_2)); + + // Get UUID-2 + let partition_after_retry = state.get_partition("data/beta").unwrap(); + let uuid_2 = partition_after_retry.uuid.clone(); + + // UUIDs should be different + assert_ne!( + uuid_1, uuid_2, + "New job should create a new UUID for the partition" + ); + + // New partition should be Building + assert_eq!( + partition_after_retry.status, + Some(crate::PartitionStatusCode::PartitionBuilding.into()), + "Retry partition should be Building" + ); + } } } }