diff --git a/features/legacy_migrator_coverage.feature b/features/legacy_migrator_coverage.feature index 9eb9ca743..4176bcb18 100644 --- a/features/legacy_migrator_coverage.feature +++ b/features/legacy_migrator_coverage.feature @@ -204,9 +204,17 @@ Feature: Legacy Data Migrator Coverage Given I have a legacy project with read-only JSON files When I run the legacy data migration Then the migration should handle backup errors gracefully + Scenario: get_all_for_project is called only once for multiple plans Given I have a legacy project with multiple plans to migrate efficiently When I run the legacy data migration with query tracking Then get_all_for_project should be called exactly once And the migration should return true And all plans should be migrated correctly + + Scenario: Map existing plan IDs correctly during migration + Given I have a legacy project with plans already in database + When I run the legacy data migration + Then existing plans should not be duplicated + And the migration should return true + And the existing plan id should be mapped without type suppression diff --git a/features/steps/legacy_migrator_steps.py b/features/steps/legacy_migrator_steps.py index 55a7bb213..e407d6474 100644 --- a/features/steps/legacy_migrator_steps.py +++ b/features/steps/legacy_migrator_steps.py @@ -1085,3 +1085,17 @@ def step_check_all_plans_migrated(context: Context) -> None: assert plan_names == expected_names, ( f"Expected plans {expected_names}, got {plan_names}" ) + + +@then("the existing plan id should be mapped without type suppression") +def step_check_existing_plan_id_mapped(context: Context) -> None: + """Check that existing plan id was mapped correctly via assert type narrowing.""" + with context.unit_of_work.transaction() as ctx: + project = ctx.projects.get_by_name("existing_plans") + assert project is not None, "Project 'existing_plans' should exist in database" + plans = ctx.plans.get_all_for_project(project.id) + main_plans = [p for p in plans if p.name == "main"] + assert len(main_plans) == 1, "Should have exactly one 'main' plan" + assert main_plans[0].id is not None, ( + "Existing plan id must be non-None (assert narrowing validates this)" + ) diff --git a/src/cleveragents/infrastructure/database/legacy_migrator.py b/src/cleveragents/infrastructure/database/legacy_migrator.py index 7e5309860..40828d48a 100644 --- a/src/cleveragents/infrastructure/database/legacy_migrator.py +++ b/src/cleveragents/infrastructure/database/legacy_migrator.py @@ -110,7 +110,8 @@ class LegacyDataMigrator: (p for p in all_plans if p.name == plan_name), None ) if existing_plan: - plan_id_map[plan_name] = existing_plan.id # type: ignore + assert existing_plan.id is not None + plan_id_map[plan_name] = existing_plan.id continue # Create new plan