From f8d332330e6853f1ab7607619cbddae595265fdd Mon Sep 17 00:00:00 2001 From: Sandeep Tuniki Date: Thu, 16 Jul 2026 18:16:22 +0530 Subject: [PATCH 1/2] Clean up dead code methods in AggregationOrchestrator and update unit tests --- .../aggregation/orchestrator.py | 30 -------------- .../aggregation/orchestrator_test.py | 41 +++++++------------ 2 files changed, 14 insertions(+), 57 deletions(-) diff --git a/pipeline/workflow/aggregation-helper/aggregation/orchestrator.py b/pipeline/workflow/aggregation-helper/aggregation/orchestrator.py index afa8fe5e..d5b9dadc 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/orchestrator.py +++ b/pipeline/workflow/aggregation-helper/aggregation/orchestrator.py @@ -335,36 +335,6 @@ def _expand_active_imports(self, active_imports: List[str]) -> List[str]: return expanded - def get_active_stages(self, active_imports: List[str]) -> List[int]: - """Returns a sorted list of unique active stage numbers across active imports.""" - expanded_imports = self._expand_active_imports(active_imports) - stages = set() - for single_import in expanded_imports: - stages.update(self._get_active_stages_for_import(single_import)) - return sorted(list(stages)) - - def execute_stage(self, stage_num: int, active_imports: List[str]) -> List[str]: - """Executes a single stage for all active imports asynchronously.""" - expanded_imports = self._expand_active_imports(active_imports) - stage_jobs = [] - for calc in self.calculations: - calc_stage = calc.get("stage", 1) - if calc_stage != stage_num: - continue - - applicable_imports = [imp for imp in expanded_imports if self._calc_applies_to_import(calc, imp)] - if not applicable_imports: - continue - - step_type = calc.get("type") - logging.info( - f"Triggering step: '{step_type}' (Stage {stage_num}) for imports {applicable_imports}..." - ) - step_jobs = self._dispatch_stage_steps(calc, applicable_imports) - stage_jobs.extend(step_jobs) - - return stage_jobs - def _execute_and_synchronize_stage(self, single_import: str, stage_num: int) -> None: """Executes a single stage for a single import and blocks until all jobs complete.""" stage_jobs = [] diff --git a/pipeline/workflow/aggregation-helper/aggregation/orchestrator_test.py b/pipeline/workflow/aggregation-helper/aggregation/orchestrator_test.py index 4083d01b..2982c932 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/orchestrator_test.py +++ b/pipeline/workflow/aggregation-helper/aggregation/orchestrator_test.py @@ -91,12 +91,12 @@ def tearDown(self): self.tmpdir.cleanup() def test_get_active_stages(self, mock_executor): - """Tests getting active stages for matching and non-matching imports.""" - stages = self.orchestrator.get_active_stages(["USFed_Census"]) - self.assertEqual(stages, [1, 2]) + """Tests resolving active stages for matching and non-matching imports via dry_run execution.""" + res_matching = self.orchestrator.run(active_imports=["USFed_Census"], dry_run=True) + self.assertEqual(res_matching.import_results["USFed_Census"].stages_executed, [1, 2]) - stages = self.orchestrator.get_active_stages(["OtherImport"]) - self.assertEqual(stages, []) + res_non_matching = self.orchestrator.run(active_imports=["OtherImport"], dry_run=True) + self.assertEqual(res_non_matching.import_results["OtherImport"].stages_executed, []) def test_directory_config_loading(self, mock_executor): """Tests that orchestrator correctly scans and loads config files from a directory.""" @@ -220,33 +220,20 @@ def test_run_skip_deletions(self, mock_entity_gen, mock_calc_gen, mock_sv_agg, m # Verify deleter was NOT called self.mock_deleter.return_value.delete_aggregated_data.assert_not_called() - def test_execute_stage(self, mock_entity_gen, mock_calc_gen, mock_sv_agg, mock_place_gen, mock_executor_cls): - """Tests manual execution of a specific stage.""" - mock_job1 = MagicMock() - mock_job1.job_id = "job-place-1" - mock_place_gen.return_value.aggregate_places.return_value = mock_job1 - - jobs = self.orchestrator.execute_stage(1, ["USFed_Census"]) - - mock_place_gen.return_value.aggregate_places.assert_called_once_with( - import_names=["USFed_Census"], - source_type="County", - destination_type="State", - allow_multiple_to_places=False - ) - self.assertEqual(jobs, [mock_job1]) - self.mock_deleter.return_value.delete_aggregated_data.assert_not_called() - - def test_execute_stage_entity_aggregation(self, mock_entity_gen, mock_calc_gen, mock_sv_agg, mock_place_gen, mock_executor_cls): - """Tests manual execution of ENTITY_AGGREGATION stage.""" + def test_run_entity_aggregation(self, mock_entity_gen, mock_calc_gen, mock_sv_agg, mock_place_gen, mock_executor_cls): + """Tests execution of ENTITY_AGGREGATION stage through orchestrator run.""" mock_job = MagicMock() mock_job.job_id = "job-entity-1" mock_entity_gen.return_value.aggregate_entities.return_value = [mock_job] - jobs = self.orchestrator.execute_stage(3, ["EarthquakeUSGS"]) - self.assertEqual(jobs, [mock_job]) + self.orchestrator.executor = MagicMock() + self.orchestrator.executor.get_jobs_status.return_value = {"status": "DONE"} + + result = self.orchestrator.run(active_imports=["EarthquakeUSGS"], dry_run=False) + self.assertTrue(result.success) + self.assertIn("EarthquakeUSGS", result.import_results) + self.assertEqual(result.import_results["EarthquakeUSGS"].stages_executed, [3]) mock_entity_gen.return_value.aggregate_entities.assert_called_once() - self.mock_deleter.return_value.delete_aggregated_data.assert_not_called() CHAINED_CONFIG_YAML = textwrap.dedent("""\ From b274138982a3b627a5f736ee1bb6f1aa5fac3bf6 Mon Sep 17 00:00:00 2001 From: Sandeep Tuniki Date: Thu, 16 Jul 2026 18:22:37 +0530 Subject: [PATCH 2/2] Apply Gemini Code Assist suggestion to strengthen entity aggregation test assertions --- .../aggregation/orchestrator_test.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/pipeline/workflow/aggregation-helper/aggregation/orchestrator_test.py b/pipeline/workflow/aggregation-helper/aggregation/orchestrator_test.py index 2982c932..ba58fe1a 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/orchestrator_test.py +++ b/pipeline/workflow/aggregation-helper/aggregation/orchestrator_test.py @@ -233,7 +233,25 @@ def test_run_entity_aggregation(self, mock_entity_gen, mock_calc_gen, mock_sv_ag self.assertTrue(result.success) self.assertIn("EarthquakeUSGS", result.import_results) self.assertEqual(result.import_results["EarthquakeUSGS"].stages_executed, [3]) + + # Verify that the entity generator was called with the correct configuration mock_entity_gen.return_value.aggregate_entities.assert_called_once() + call_args = mock_entity_gen.return_value.aggregate_entities.call_args[0][0] + self.assertEqual(len(call_args), 1) + config = call_args[0] + self.assertEqual(config.entity_types, ["EarthquakeEvent"]) + self.assertEqual(config.location_props, ["affectedPlace"]) + self.assertEqual(config.date_prop, "occurrenceTime") + self.assertEqual(config.agg_date_formats, ["YYYY"]) + self.assertEqual(config.constraints, [{"property": "magnitude", "min": 7, "unit": "M"}]) + self.assertEqual(config.output_import, "EarthquakeUSGS_Agg") + self.assertEqual(config.input_imports, ["EarthquakeUSGS"]) + + # Verify deleter was called with expected outputs + self.mock_deleter.return_value.delete_aggregated_data.assert_called_once_with( + ["EarthquakeUSGS_Agg"] + ) + CHAINED_CONFIG_YAML = textwrap.dedent("""\