From 19b4a34f1711d51da15ae7d520185932e53d3fa9 Mon Sep 17 00:00:00 2001 From: Jvst Me Date: Thu, 20 Aug 2026 00:51:20 +0200 Subject: [PATCH] Make `GatewayReplicaModel.backend_id` required The main motivation of this change is to ensure that no legacy Sky gateways with `backend_id=None` remain in any databases, as well as simplify reasoning about the data model, since `backend_id` isn't actually set to `None` anywhere as of this version. --- .../pipeline_tasks/gateway_replicas.py | 6 - ..._require_gatewayreplicamodel_backend_id.py | 50 +++++++ src/dstack/_internal/server/models.py | 6 +- .../server/services/gateways/__init__.py | 4 +- src/dstack/_internal/server/testing/common.py | 11 +- .../pipeline_tasks/test_gateway_replicas.py | 105 ++++++++------ .../pipeline_tasks/test_gateways.py | 129 ++++++++++++++---- .../pipeline_tasks/test_running_jobs.py | 25 ++-- .../pipeline_tasks/test_runs/test_active.py | 25 ++-- .../pipeline_tasks/test_runs/test_pending.py | 4 +- .../_internal/server/routers/test_gateways.py | 76 +++++------ .../_internal/server/routers/test_runs.py | 18 +-- .../server/services/gateways/test_gateways.py | 4 +- 13 files changed, 303 insertions(+), 160 deletions(-) create mode 100644 src/dstack/_internal/server/migrations/versions/2026/08_18_1051_eee3e79f29e9_require_gatewayreplicamodel_backend_id.py diff --git a/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py b/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py index 6cf455812..5e78a1352 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py @@ -452,8 +452,6 @@ async def _provision_gateway_replica( replica_model: GatewayReplicaModel, ) -> _GatewayReplicaUpdateMap: try: - if replica_model.backend_id is None: # unexpected - raise BackendNotAvailable() (_, backend) = await backends_services.get_project_backend_with_model_by_id_or_error( project=gateway_model.project, backend_id=replica_model.backend_id ) @@ -614,8 +612,6 @@ async def _register_replica_with_load_balancer( if replica_model.instance_id is None: return "instance_id is None, cannot register with load balancer" try: - if replica_model.backend_id is None: - raise BackendNotAvailable() (_, backend) = await backends_services.get_project_backend_with_model_by_id_or_error( project=gateway_model.project, backend_id=replica_model.backend_id ) @@ -1565,8 +1561,6 @@ async def _process_terminating_item(item: GatewayReplicaPipelineItem): status=GatewayReplicaStatus.TERMINATED, active=False, deleted=True ) try: - if replica_model.backend_id is None: # unexpected - raise BackendNotAvailable() (_, backend) = await backends_services.get_project_backend_with_model_by_id_or_error( project=gateway_model.project, backend_id=replica_model.backend_id, diff --git a/src/dstack/_internal/server/migrations/versions/2026/08_18_1051_eee3e79f29e9_require_gatewayreplicamodel_backend_id.py b/src/dstack/_internal/server/migrations/versions/2026/08_18_1051_eee3e79f29e9_require_gatewayreplicamodel_backend_id.py new file mode 100644 index 000000000..7aa3626a5 --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/2026/08_18_1051_eee3e79f29e9_require_gatewayreplicamodel_backend_id.py @@ -0,0 +1,50 @@ +"""Require GatewayReplicaModel.backend_id + +Revision ID: eee3e79f29e9 +Revises: 04126c7ea0c8 +Create Date: 2026-08-18 10:51:19.731160+00:00 + +""" + +import sqlalchemy as sa +import sqlalchemy_utils +from alembic import op + +# revision identifiers, used by Alembic. +revision = "eee3e79f29e9" +down_revision = "04126c7ea0c8" +branch_labels = None +depends_on = None + +# partial definition for queries +gateway_computes = sa.table( + "gateway_computes", + sa.column("id"), + sa.column("backend_id"), +) + + +def upgrade() -> None: + # Gateway computes with backend_id=None were only possible in dstack Sky and were already + # removed from all significant environments before this migration. So this cleanup query is + # expected to be a NOOP in most cases, except for, possibly, some local dev databases. + op.execute(sa.delete(gateway_computes).where(gateway_computes.c.backend_id.is_(None))) + + with op.batch_alter_table("gateway_computes", schema=None) as batch_op: + batch_op.alter_column( + "backend_id", + existing_type=sqlalchemy_utils.types.uuid.UUIDType(binary=False), + nullable=False, + ) + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("gateway_computes", schema=None) as batch_op: + batch_op.alter_column( + "backend_id", + existing_type=sqlalchemy_utils.types.uuid.UUIDType(binary=False), + nullable=True, + ) + + # ### end Alembic commands ### diff --git a/src/dstack/_internal/server/models.py b/src/dstack/_internal/server/models.py index dcc5707f4..24822d7c8 100644 --- a/src/dstack/_internal/server/models.py +++ b/src/dstack/_internal/server/models.py @@ -741,10 +741,8 @@ class GatewayReplicaModel(PipelineModelMixin, BaseModel): Use `gateway or legacy_gateway` to get the gateway regardless of version. """ - backend_id: Mapped[Optional[uuid.UUID]] = mapped_column( - ForeignKey("backends.id", ondelete="CASCADE") - ) - backend: Mapped[Optional["BackendModel"]] = relationship() + backend_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("backends.id", ondelete="CASCADE")) + backend: Mapped["BackendModel"] = relationship() ssh_private_key: Mapped[str] = mapped_column(Text) """`ssh_private_key` is the key used to authorize the server with the gateway.""" diff --git a/src/dstack/_internal/server/services/gateways/__init__.py b/src/dstack/_internal/server/services/gateways/__init__.py index a9fec7aca..acaa890cb 100644 --- a/src/dstack/_internal/server/services/gateways/__init__.py +++ b/src/dstack/_internal/server/services/gateways/__init__.py @@ -948,7 +948,7 @@ def gateway_model_to_gateway( all_replica_models = sorted( get_gateway_replica_models(gateway_model), key=lambda r: r.replica_num ) - relevant_replica_models = [] + relevant_replica_models: list[GatewayReplicaModel] = [] for replica_num, replica_models_for_num in itertools.groupby( all_replica_models, key=lambda r: r.replica_num ): @@ -959,7 +959,7 @@ def gateway_model_to_gateway( GatewayReplica( hostname=replica_model.ip_address, replica_num=replica_model.replica_num, - backend=replica_model.backend.type if replica_model.backend else None, + backend=replica_model.backend.type, region=replica_model.region, created_at=replica_model.created_at, status=replica_model.status, diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py index 0fa5b3556..0f11da4a2 100644 --- a/src/dstack/_internal/server/testing/common.py +++ b/src/dstack/_internal/server/testing/common.py @@ -711,8 +711,8 @@ async def create_gateway( async def create_gateway_replica( session: AsyncSession, + backend: BackendModel, gateway_id: Optional[UUID] = None, - backend_id: Optional[UUID] = None, ip_address: Optional[str] = "1.1.1.1", region: Optional[str] = "us", instance_id: Optional[str] = "i-1234567890", @@ -734,16 +734,11 @@ async def create_gateway_replica( testing against both in major test cases. """ if configuration is None and populate_configuration: - backend_type = BackendType.AWS - if backend_id is not None: - backend = await session.get(BackendModel, backend_id) - assert backend is not None - backend_type = backend.type assert region is not None configuration = GatewayReplicaConfiguration( project_name="test-project", instance_name=instance_id or "test-instance", - backend=backend_type, + backend=backend.type, region=region, public_ip=True, ssh_key_pub=ssh_public_key, @@ -751,7 +746,7 @@ async def create_gateway_replica( ).model_dump_json() gateway_replica = GatewayReplicaModel( gateway_id=gateway_id, - backend_id=backend_id, + backend_id=backend.id, ip_address=ip_address, region=region, instance_id=instance_id, diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_gateway_replicas.py b/src/tests/_internal/server/background/pipeline_tasks/test_gateway_replicas.py index 9c2e0356e..15128041b 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_gateway_replicas.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_gateway_replicas.py @@ -109,6 +109,7 @@ async def test_fetch_selects_eligible_replicas_and_sets_lock_fields( submitted = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, ip_address=None, instance_id=None, region=None, @@ -119,12 +120,14 @@ async def test_fetch_selects_eligible_replicas_and_sets_lock_fields( provisioning = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, status=GatewayReplicaStatus.PROVISIONING, last_processed_at=stale - timedelta(seconds=2), ) terminating = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, status=GatewayReplicaStatus.TERMINATING, active=False, last_processed_at=stale - timedelta(seconds=1), @@ -132,12 +135,14 @@ async def test_fetch_selects_eligible_replicas_and_sets_lock_fields( running = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, last_processed_at=stale, ) terminated = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, status=GatewayReplicaStatus.TERMINATED, active=False, last_processed_at=stale, @@ -145,6 +150,7 @@ async def test_fetch_selects_eligible_replicas_and_sets_lock_fields( recent = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, status=GatewayReplicaStatus.SUBMITTED, ip_address=None, instance_id=None, @@ -157,6 +163,7 @@ async def test_fetch_selects_eligible_replicas_and_sets_lock_fields( locked = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, status=GatewayReplicaStatus.SUBMITTED, ip_address=None, instance_id=None, @@ -212,6 +219,7 @@ async def test_fetch_includes_recent_replica_with_skip_min_processing_interval( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, last_processed_at=now, ) @@ -254,6 +262,7 @@ async def test_fetch_includes_running_replica_needing_cleanup( if legacy_replica: replica = await create_gateway_replica( session=session, + backend=backend, status=GatewayReplicaStatus.RUNNING, last_processed_at=stale, ) @@ -262,6 +271,7 @@ async def test_fetch_includes_running_replica_needing_cleanup( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, last_processed_at=stale, ) @@ -281,10 +291,13 @@ async def test_fetch_includes_running_replica_with_hard_deleted_gateway( ): # A replica whose gateway was hard-deleted (orphaned). The fetcher should # pick it up so the worker can log the error. + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) stale = get_current_datetime() - timedelta(minutes=1) replica = await create_gateway_replica( session=session, gateway_id=None, + backend=backend, status=GatewayReplicaStatus.RUNNING, last_processed_at=stale, ) @@ -318,6 +331,7 @@ async def test_fetch_includes_running_replica_with_healthy_gateway( if legacy_replica: replica = await create_gateway_replica( session=session, + backend=backend, status=GatewayReplicaStatus.RUNNING, last_processed_at=stale, ) @@ -326,6 +340,7 @@ async def test_fetch_includes_running_replica_with_healthy_gateway( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, last_processed_at=stale, ) @@ -355,6 +370,7 @@ async def test_fetch_includes_running_replica_marked_for_scale_in( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, last_processed_at=stale, ) @@ -385,7 +401,7 @@ async def test_submitted_to_provisioning( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, ip_address=None, instance_id=None, region=None, @@ -431,7 +447,7 @@ async def test_submitted_backend_error_marks_terminated( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, ip_address=None, instance_id=None, region=None, @@ -471,7 +487,7 @@ async def test_submitted_backend_not_available_marks_terminated( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, ip_address=None, instance_id=None, region=None, @@ -507,7 +523,7 @@ async def test_submitted_skips_provisioning_if_gateway_to_be_deleted( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, ip_address=None, instance_id=None, region=None, @@ -542,7 +558,7 @@ async def test_submitted_skips_provisioning_if_gateway_failed( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, ip_address=None, instance_id=None, region=None, @@ -577,7 +593,7 @@ async def test_submitted_unexpected_error_marks_terminated( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, ip_address=None, instance_id=None, region=None, @@ -618,7 +634,7 @@ async def test_submitted_to_terminated_when_scaled_in( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, ip_address=None, instance_id=None, region=None, @@ -677,6 +693,7 @@ async def test_running_to_terminating( if legacy_replica: replica = await create_gateway_replica( session=session, + backend=backend, status=GatewayReplicaStatus.RUNNING, active=True, populate_configuration=populate_configuration, @@ -686,6 +703,7 @@ async def test_running_to_terminating( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, active=True, populate_configuration=populate_configuration, @@ -713,6 +731,7 @@ async def test_running_to_terminating_when_scaled_in( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, active=True, ) @@ -812,7 +831,7 @@ async def test_registers_new_service_and_replica( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ssh_private_key="replica-private-key", ) @@ -906,7 +925,7 @@ async def test_unregisters_dangling_service_and_stale_replica( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ) # A live, still-expected service with one live replica and one stale @@ -1020,7 +1039,7 @@ async def test_deletes_registration_models_for_unregistered_service_and_replica( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ) # A live, still-expected service with one live replica and one stale @@ -1165,7 +1184,7 @@ async def test_unregisters_replicas_of_dangling_service_without_extra_gateway_ca replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ) # A finished run whose service and replica are still (erroneously) @@ -1255,7 +1274,7 @@ async def test_no_gateway_calls_when_state_already_in_sync( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ) run, job = await self._create_service_run_and_job( @@ -1326,7 +1345,7 @@ async def test_recovers_legacy_service_id_by_matching_replica( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ) run, job = await self._create_service_run_and_job( @@ -1376,7 +1395,7 @@ async def test_unregisters_and_reregisters_legacy_service_without_id_and_replica replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ) run, job = await self._create_service_run_and_job( @@ -1476,7 +1495,7 @@ async def test_does_nothing_when_in_sync_and_registrations_already_exist( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ) run, job = await self._create_service_run_and_job( @@ -1570,7 +1589,7 @@ async def test_reconciles_out_of_sync_registration_models( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ) run, job = await self._create_service_run_and_job( @@ -1673,7 +1692,7 @@ async def test_propagates_registration_errors_and_increments_register_attempt( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ) run_replica_fails, job_replica_fails = await self._create_service_run_and_job( @@ -1805,7 +1824,7 @@ async def test_does_not_emit_duplicate_registration_error_event_for_unchanged_er replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ) run, job = await self._create_service_run_and_job( @@ -1895,7 +1914,7 @@ async def test_propagates_unregistration_errors_and_increments_unregister_attemp replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ) # A live, still-expected service with one live replica and one stale @@ -2008,7 +2027,7 @@ async def test_does_not_emit_duplicate_unregistration_error_event_for_unchanged_ replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ) run, job = await self._create_service_run_and_job( @@ -2101,6 +2120,7 @@ async def test_provisioning_to_running( if legacy_replica: replica = await create_gateway_replica( session=session, + backend=backend, status=GatewayReplicaStatus.PROVISIONING, populate_configuration=populate_configuration, ) @@ -2109,6 +2129,7 @@ async def test_provisioning_to_running( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, status=GatewayReplicaStatus.PROVISIONING, populate_configuration=populate_configuration, ) @@ -2144,7 +2165,7 @@ async def test_provisioning_to_running_registers_with_load_balancer( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.PROVISIONING, ) _lock_replica(replica) @@ -2193,7 +2214,7 @@ async def test_provisioning_skips_load_balancer_registration_without_hostname( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.PROVISIONING, ) _lock_replica(replica) @@ -2235,7 +2256,7 @@ async def test_provisioning_to_terminating_when_load_balancer_registration_fails replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.PROVISIONING, ) _lock_replica(replica) @@ -2282,7 +2303,7 @@ async def test_provisioning_to_terminating_when_backend_does_not_support_load_ba replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.PROVISIONING, ) _lock_replica(replica) @@ -2325,7 +2346,7 @@ async def test_provisioning_waits_for_pending_acm_gateway_migration( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.PROVISIONING, hostname_deprecated_readonly="legacy-lb.example.com", ) @@ -2359,6 +2380,7 @@ async def test_provisioning_to_terminating_if_connect_fails( if legacy_replica: replica = await create_gateway_replica( session=session, + backend=backend, status=GatewayReplicaStatus.PROVISIONING, ) gateway.gateway_replica_id = replica.id @@ -2366,6 +2388,7 @@ async def test_provisioning_to_terminating_if_connect_fails( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, status=GatewayReplicaStatus.PROVISIONING, ) _lock_replica(replica) @@ -2398,6 +2421,7 @@ async def test_provisioning_to_terminating_if_configure_fails( if legacy_replica: replica = await create_gateway_replica( session=session, + backend=backend, status=GatewayReplicaStatus.PROVISIONING, ) gateway.gateway_replica_id = replica.id @@ -2405,6 +2429,7 @@ async def test_provisioning_to_terminating_if_configure_fails( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, status=GatewayReplicaStatus.PROVISIONING, ) _lock_replica(replica) @@ -2458,6 +2483,7 @@ async def test_provisioning_to_terminating_if_gateway_needs_cleanup( if legacy_replica: replica = await create_gateway_replica( session=session, + backend=backend, status=GatewayReplicaStatus.PROVISIONING, ) gateway.gateway_replica_id = replica.id @@ -2465,6 +2491,7 @@ async def test_provisioning_to_terminating_if_gateway_needs_cleanup( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, status=GatewayReplicaStatus.PROVISIONING, ) _lock_replica(replica) @@ -2506,7 +2533,7 @@ async def test_terminating_to_terminated( if legacy_replica: replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.TERMINATING, active=False, populate_configuration=populate_configuration, @@ -2516,7 +2543,7 @@ async def test_terminating_to_terminated( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.TERMINATING, active=False, populate_configuration=populate_configuration, @@ -2563,7 +2590,7 @@ async def test_terminating_to_terminated_deletes_only_own_registration_records( replica_to_terminate = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.TERMINATING, active=False, replica_num=0, @@ -2571,7 +2598,7 @@ async def test_terminating_to_terminated_deletes_only_own_registration_records( other_replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ip_address="2.2.2.2", instance_id="i-eeeeeeeeee", @@ -2646,7 +2673,7 @@ async def test_terminating_deregisters_from_load_balancer_before_terminating( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.TERMINATING, active=False, ) @@ -2699,7 +2726,7 @@ async def test_terminating_proceeds_when_load_balancer_deregistration_raises( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.TERMINATING, active=False, ) @@ -2750,7 +2777,7 @@ async def test_terminating_skips_deregistration_when_gateway_has_no_hostname( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.TERMINATING, active=False, ) @@ -2793,7 +2820,7 @@ async def test_terminating_waits_for_pending_acm_gateway_migration( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.TERMINATING, active=False, hostname_deprecated_readonly="legacy-lb.example.com", @@ -2834,7 +2861,7 @@ async def test_terminating_to_terminated_if_backend_not_available( if legacy_replica: replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.TERMINATING, active=False, ) @@ -2843,7 +2870,7 @@ async def test_terminating_to_terminated_if_backend_not_available( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.TERMINATING, active=False, ) @@ -2876,7 +2903,7 @@ async def test_terminating_to_terminated_with_no_instance_id( if legacy_replica: replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, instance_id=None, status=GatewayReplicaStatus.TERMINATING, active=False, @@ -2886,7 +2913,7 @@ async def test_terminating_to_terminated_with_no_instance_id( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, instance_id=None, status=GatewayReplicaStatus.TERMINATING, active=False, @@ -2931,7 +2958,7 @@ async def test_terminating_retries_if_terminate_fails( if legacy_replica: replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.TERMINATING, active=False, ) @@ -2940,7 +2967,7 @@ async def test_terminating_retries_if_terminate_fails( replica = await create_gateway_replica( session=session, gateway_id=gateway.id, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.TERMINATING, active=False, ) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_gateways.py b/src/tests/_internal/server/background/pipeline_tasks/test_gateways.py index e2d38f651..cc9aad841 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_gateways.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_gateways.py @@ -277,13 +277,14 @@ async def test_fetch_excludes_running_gateway_when_replica_count_matches( if legacy_replica: replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ) gateway.gateway_replica_id = replica.id else: await create_gateway_replica( session=session, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, ) @@ -309,6 +310,7 @@ async def test_fetch_includes_running_gateway_with_pending_scale_attempt_even_if ) await create_gateway_replica( session=session, + backend=backend, gateway_id=gateway.id, ip_address=None, instance_id=None, @@ -340,13 +342,14 @@ async def test_fetch_includes_running_gateway_when_replica_count_not_matches( if legacy_replica: replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ) gateway.gateway_replica_id = replica.id else: await create_gateway_replica( session=session, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, ) @@ -375,7 +378,7 @@ async def test_fetch_includes_running_gateway_with_unmigrated_legacy_hostname( if legacy_replica: replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, hostname_deprecated_readonly="legacy-lb.example.com", ) @@ -383,6 +386,7 @@ async def test_fetch_includes_running_gateway_with_unmigrated_legacy_hostname( else: await create_gateway_replica( session=session, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, hostname_deprecated_readonly="legacy-lb.example.com", @@ -412,6 +416,7 @@ async def test_fetch_excludes_running_gateway_without_legacy_hostname_to_migrate ) await create_gateway_replica( session=session, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, hostname_deprecated_readonly=None, @@ -439,6 +444,7 @@ async def test_fetch_excludes_already_migrated_gateway( ) await create_gateway_replica( session=session, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, hostname_deprecated_readonly="legacy-lb.example.com", @@ -652,7 +658,7 @@ async def test_provisioning_to_running( if legacy_replica: gateway_replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, populate_configuration=populate_configuration, ) @@ -660,6 +666,7 @@ async def test_provisioning_to_running( else: await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, populate_configuration=populate_configuration, @@ -691,6 +698,7 @@ async def test_provisioning_migrates_hostname_and_backend_data_from_legacy_repli ) await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, hostname_deprecated_readonly="legacy-lb.example.com", @@ -723,6 +731,7 @@ async def test_provisioning_does_not_overwrite_already_migrated_hostname( ) await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, hostname_deprecated_readonly="stale-legacy-lb.example.com", @@ -753,6 +762,7 @@ async def test_provisioning_to_running_with_multiple_replicas( ) await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, ip_address="1.1.1.1", status=GatewayReplicaStatus.RUNNING, @@ -760,6 +770,7 @@ async def test_provisioning_to_running_with_multiple_replicas( ) await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, ip_address="2.2.2.2", status=GatewayReplicaStatus.RUNNING, @@ -791,6 +802,7 @@ async def test_still_provisioning_if_not_all_replicas_running( ) await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, ip_address="1.1.1.1", status=GatewayReplicaStatus.RUNNING, @@ -798,6 +810,7 @@ async def test_still_provisioning_if_not_all_replicas_running( ) await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, ip_address="2.2.2.2", status=GatewayReplicaStatus.PROVISIONING, @@ -839,7 +852,7 @@ async def test_marks_gateway_as_failed_if_replica_failed( if legacy_replica: gateway_replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, status=replica_status, active=False, ) @@ -847,6 +860,7 @@ async def test_marks_gateway_as_failed_if_replica_failed( else: await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, status=replica_status, active=False, @@ -880,6 +894,7 @@ async def test_still_provisioning_with_submitted_replica( ) await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, ip_address=None, instance_id=None, @@ -923,7 +938,7 @@ async def test_still_provisioning_when_scale_out_adds_new_replicas( if legacy_replica: gateway_replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, ip_address="1.1.1.1", status=GatewayReplicaStatus.RUNNING, populate_configuration=populate_configuration, @@ -932,6 +947,7 @@ async def test_still_provisioning_when_scale_out_adds_new_replicas( else: await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, ip_address="1.1.1.1", status=GatewayReplicaStatus.RUNNING, @@ -970,6 +986,7 @@ async def test_provisioning_to_running_when_scale_in_removes_surplus_replicas( ) older = await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, ip_address="1.1.1.1", status=GatewayReplicaStatus.PROVISIONING, @@ -978,6 +995,7 @@ async def test_provisioning_to_running_when_scale_in_removes_surplus_replicas( older.created_at = datetime(2025, 1, 1) newer = await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, ip_address="2.2.2.2", status=GatewayReplicaStatus.RUNNING, @@ -1014,6 +1032,7 @@ async def test_ignores_previously_scaled_in_replica_when_determining_status( ) await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, ip_address="1.1.1.1", status=GatewayReplicaStatus.RUNNING, @@ -1021,6 +1040,7 @@ async def test_ignores_previously_scaled_in_replica_when_determining_status( ) scaled_in_replica = await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, ip_address="2.2.2.2", status=GatewayReplicaStatus.TERMINATING, @@ -1060,13 +1080,14 @@ async def test_no_scaling_when_replica_count_matches( if legacy_replica: replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ) gateway.gateway_replica_id = replica.id else: await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0, @@ -1103,6 +1124,7 @@ async def test_running_migrates_hostname_and_backend_data_from_legacy_replica( ) await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, hostname_deprecated_readonly="legacy-lb.example.com", @@ -1141,7 +1163,7 @@ async def test_scales_out_when_desired_replica_count_increased( if legacy_replica: gateway_replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, populate_configuration=populate_configuration, ) @@ -1149,6 +1171,7 @@ async def test_scales_out_when_desired_replica_count_increased( else: await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0, @@ -1187,15 +1210,27 @@ async def test_scales_in_oldest_replicas_when_desired_replica_count_decreased( replicas=1, ) replica0 = await create_gateway_replica( - session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 + session, + backend=backend, + gateway_id=gateway.id, + status=GatewayReplicaStatus.RUNNING, + replica_num=0, ) replica0.created_at = datetime(2025, 1, 1) replica1 = await create_gateway_replica( - session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=1 + session, + backend=backend, + gateway_id=gateway.id, + status=GatewayReplicaStatus.RUNNING, + replica_num=1, ) replica1.created_at = datetime(2025, 1, 2) replica2 = await create_gateway_replica( - session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=2 + session, + backend=backend, + gateway_id=gateway.id, + status=GatewayReplicaStatus.RUNNING, + replica_num=2, ) replica2.created_at = datetime(2025, 1, 3) gateway.lock_token = uuid.uuid4() @@ -1226,11 +1261,16 @@ async def test_scale_in_prefers_less_advanced_replicas_over_older_running_ones( replicas=1, ) running = await create_gateway_replica( - session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 + session, + backend=backend, + gateway_id=gateway.id, + status=GatewayReplicaStatus.RUNNING, + replica_num=0, ) running.created_at = datetime(2025, 1, 1) submitted = await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.SUBMITTED, replica_num=1, @@ -1266,13 +1306,17 @@ async def test_no_scaling_for_legacy_gateway_without_desired_replica_count( if legacy_replica: replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.RUNNING, ) gateway.gateway_replica_id = replica.id else: await create_gateway_replica( - session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 + session, + backend=backend, + gateway_id=gateway.id, + status=GatewayReplicaStatus.RUNNING, + replica_num=0, ) gateway.desired_replica_count = None gateway.lock_token = uuid.uuid4() @@ -1300,7 +1344,11 @@ async def test_scale_out_skipped_before_retry_delay_elapses( replicas=2, ) await create_gateway_replica( - session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 + session, + backend=backend, + gateway_id=gateway.id, + status=GatewayReplicaStatus.RUNNING, + replica_num=0, ) gateway.replica_scale_attempt = 1 gateway.last_replica_scale_attempt_at = get_current_datetime() @@ -1328,7 +1376,11 @@ async def test_scale_out_retries_after_retry_delay_elapses( replicas=2, ) await create_gateway_replica( - session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 + session, + backend=backend, + gateway_id=gateway.id, + status=GatewayReplicaStatus.RUNNING, + replica_num=0, ) gateway.replica_scale_attempt = 1 gateway.last_replica_scale_attempt_at = get_current_datetime() - timedelta(minutes=5) @@ -1356,7 +1408,11 @@ async def test_scale_out_stops_after_reaching_attempt_limit( replicas=2, ) await create_gateway_replica( - session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 + session, + backend=backend, + gateway_id=gateway.id, + status=GatewayReplicaStatus.RUNNING, + replica_num=0, ) gateway.replica_scale_attempt = _MAX_REPLICA_SCALE_ATTEMPTS gateway.last_replica_scale_attempt_at = datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc) @@ -1386,7 +1442,11 @@ async def test_scale_out_emits_event_on_reaching_attempt_limit( replicas=2, ) await create_gateway_replica( - session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 + session, + backend=backend, + gateway_id=gateway.id, + status=GatewayReplicaStatus.RUNNING, + replica_num=0, ) gateway.replica_scale_attempt = _MAX_REPLICA_SCALE_ATTEMPTS - 1 gateway.last_replica_scale_attempt_at = datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc) @@ -1419,6 +1479,7 @@ async def test_attempt_counter_not_reset_while_replacement_replica_still_provisi ) await create_gateway_replica( session, + backend=backend, gateway_id=gateway.id, ip_address=None, instance_id=None, @@ -1450,7 +1511,11 @@ async def test_attempt_counter_resets_and_scales_out_immediately_after_in_place_ replicas=2, ) await create_gateway_replica( - session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 + session, + backend=backend, + gateway_id=gateway.id, + status=GatewayReplicaStatus.RUNNING, + replica_num=0, ) gateway.replica_scale_attempt = _MAX_REPLICA_SCALE_ATTEMPTS gateway.last_replica_scale_attempt_at = datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc) @@ -1480,7 +1545,11 @@ async def test_attempt_counter_not_reset_when_update_precedes_last_scale_attempt replicas=2, ) await create_gateway_replica( - session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 + session, + backend=backend, + gateway_id=gateway.id, + status=GatewayReplicaStatus.RUNNING, + replica_num=0, ) gateway.replica_scale_attempt = _MAX_REPLICA_SCALE_ATTEMPTS gateway.last_replica_scale_attempt_at = get_current_datetime() @@ -1546,7 +1615,7 @@ async def test_deletes_gateway_when_all_replicas_terminated( if legacy_replica: gateway_replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, status=GatewayReplicaStatus.TERMINATED, active=False, populate_configuration=populate_configuration, @@ -1555,7 +1624,7 @@ async def test_deletes_gateway_when_all_replicas_terminated( else: await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.TERMINATED, active=False, @@ -1590,7 +1659,7 @@ async def test_deletes_gateway_and_terminates_load_balancer_when_hostname_set( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.TERMINATED, active=False, @@ -1634,7 +1703,7 @@ async def test_delete_skips_load_balancer_termination_when_hostname_not_set( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.TERMINATED, active=False, @@ -1669,7 +1738,7 @@ async def test_delete_deferred_when_load_balancer_termination_fails( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.TERMINATED, active=False, @@ -1719,7 +1788,7 @@ async def test_delete_deferred_when_backend_does_not_support_load_balancer( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.TERMINATED, active=False, @@ -1757,7 +1826,7 @@ async def test_delete_migrates_hostname_before_evaluating_termination( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, status=GatewayReplicaStatus.TERMINATED, active=False, @@ -1810,7 +1879,7 @@ async def test_waits_when_replicas_not_yet_terminated( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, status=replica_status, active=False, @@ -1849,7 +1918,7 @@ async def test_deletes_gateway_with_multiple_replicas_all_terminated( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, ip_address="1.1.1.1", status=GatewayReplicaStatus.TERMINATED, @@ -1858,7 +1927,7 @@ async def test_deletes_gateway_with_multiple_replicas_all_terminated( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, ip_address="2.2.2.2", status=GatewayReplicaStatus.TERMINATED, diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py index 560e3c201..e79363422 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py @@ -1888,7 +1888,9 @@ async def test_terminates_job_on_gateway_registration_failure( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - gateway_replica = await create_gateway_replica(session=session, gateway_id=gateway.id) + gateway_replica = await create_gateway_replica( + session=session, gateway_id=gateway.id, backend=backend + ) run = await create_run( session=session, project=project, @@ -1967,10 +1969,10 @@ async def test_does_not_terminate_job_when_registered_with_at_least_one_gateway_ status=GatewayStatus.RUNNING, ) gateway_replica_1 = await create_gateway_replica( - session=session, gateway_id=gateway.id, replica_num=0 + session=session, gateway_id=gateway.id, backend=backend, replica_num=0 ) gateway_replica_2 = await create_gateway_replica( - session=session, gateway_id=gateway.id, replica_num=1 + session=session, gateway_id=gateway.id, backend=backend, replica_num=1 ) run = await create_run( session=session, @@ -2058,10 +2060,12 @@ async def test_does_not_terminate_job_when_gateway_replica_has_not_attempted_reg status=GatewayStatus.RUNNING, ) gateway_replica_1 = await create_gateway_replica( - session=session, gateway_id=gateway.id, replica_num=0 + session=session, gateway_id=gateway.id, backend=backend, replica_num=0 ) # Second running replica has not attempted registration yet (e.g. just came up). - await create_gateway_replica(session=session, gateway_id=gateway.id, replica_num=1) + await create_gateway_replica( + session=session, gateway_id=gateway.id, backend=backend, replica_num=1 + ) run = await create_run( session=session, project=project, @@ -2140,13 +2144,14 @@ async def test_terminates_job_ignoring_registration_on_non_running_replica( status=GatewayStatus.RUNNING, ) gateway_replica_running = await create_gateway_replica( - session=session, gateway_id=gateway.id, replica_num=0 + session=session, gateway_id=gateway.id, backend=backend, replica_num=0 ) # Terminated replica successfully registered before going away — should be ignored, # since only currently running replicas count towards the predicate. gateway_replica_terminating = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, replica_num=1, status=GatewayReplicaStatus.TERMINATING, ) @@ -2447,14 +2452,14 @@ async def test_registers_service_replica_in_gateway( if legacy_replica: gateway_replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, ) gateway.gateway_replica_id = gateway_replica.id await session.commit() else: gateway_replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, ) run = await create_run( @@ -2540,7 +2545,7 @@ async def test_registers_service_replica_in_gateway_when_running_on_imported_ins ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, ) run = await create_run( @@ -2829,7 +2834,7 @@ async def test_registers_router_replica_but_not_worker_replica_in_gateway( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, ) run = await create_run( diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py b/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py index fe24a8cde..389f70f45 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py @@ -196,7 +196,9 @@ async def test_terminates_run_on_gateway_registration_failure( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - gateway_replica = await create_gateway_replica(session=session, gateway_id=gateway.id) + gateway_replica = await create_gateway_replica( + session=session, gateway_id=gateway.id, backend=backend + ) run = await create_run( session=session, project=project, @@ -250,10 +252,10 @@ async def test_does_not_terminate_run_when_service_registered_despite_failed_att status=GatewayStatus.RUNNING, ) gateway_replica_1 = await create_gateway_replica( - session=session, gateway_id=gateway.id, replica_num=0 + session=session, gateway_id=gateway.id, backend=backend, replica_num=0 ) gateway_replica_2 = await create_gateway_replica( - session=session, gateway_id=gateway.id, replica_num=1 + session=session, gateway_id=gateway.id, backend=backend, replica_num=1 ) run = await create_run( session=session, @@ -316,10 +318,12 @@ async def test_does_not_terminate_run_when_one_running_replica_has_not_attempted status=GatewayStatus.RUNNING, ) gateway_replica_1 = await create_gateway_replica( - session=session, gateway_id=gateway.id, replica_num=0 + session=session, gateway_id=gateway.id, backend=backend, replica_num=0 ) # Second running replica has not attempted registration yet (e.g. just came up). - await create_gateway_replica(session=session, gateway_id=gateway.id, replica_num=1) + await create_gateway_replica( + session=session, gateway_id=gateway.id, backend=backend, replica_num=1 + ) run = await create_run( session=session, project=project, @@ -373,13 +377,14 @@ async def test_terminates_run_ignoring_registration_on_non_running_replica( status=GatewayStatus.RUNNING, ) gateway_replica_running = await create_gateway_replica( - session=session, gateway_id=gateway.id, replica_num=0 + session=session, gateway_id=gateway.id, backend=backend, replica_num=0 ) # Terminated replica successfully registered before going away — should be ignored, # since only currently running replicas count towards the predicate. gateway_replica_terminating = await create_gateway_replica( session=session, gateway_id=gateway.id, + backend=backend, replica_num=1, status=GatewayReplicaStatus.TERMINATING, ) @@ -1408,10 +1413,10 @@ async def test_service_rolling_deployment_keeps_old_replica_until_new_replica_re status=GatewayStatus.RUNNING, ) gateway_replica_1 = await create_gateway_replica( - session=session, gateway_id=gateway.id, replica_num=0 + session=session, gateway_id=gateway.id, backend=backend, replica_num=0 ) gateway_replica_2 = await create_gateway_replica( - session=session, gateway_id=gateway.id, replica_num=1 + session=session, gateway_id=gateway.id, backend=backend, replica_num=1 ) run_spec = get_run_spec( repo_id=repo.name, @@ -1518,10 +1523,10 @@ async def test_service_rolling_deployment_scales_down_old_replica_once_new_repli status=GatewayStatus.RUNNING, ) gateway_replica_1 = await create_gateway_replica( - session=session, gateway_id=gateway.id, replica_num=0 + session=session, gateway_id=gateway.id, backend=backend, replica_num=0 ) gateway_replica_2 = await create_gateway_replica( - session=session, gateway_id=gateway.id, replica_num=1 + session=session, gateway_id=gateway.id, backend=backend, replica_num=1 ) run_spec = get_run_spec( repo_id=repo.name, diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pending.py b/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pending.py index af41ba839..6cb006821 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pending.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pending.py @@ -381,7 +381,9 @@ async def test_terminates_run_on_gateway_registration_failure( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - gateway_replica = await create_gateway_replica(session=session, gateway_id=gateway.id) + gateway_replica = await create_gateway_replica( + session=session, gateway_id=gateway.id, backend=backend + ) run_spec = get_run_spec( run_name="test-run", repo_id=repo.name, diff --git a/src/tests/_internal/server/routers/test_gateways.py b/src/tests/_internal/server/routers/test_gateways.py index 0ff8e2ca5..667c01714 100644 --- a/src/tests/_internal/server/routers/test_gateways.py +++ b/src/tests/_internal/server/routers/test_gateways.py @@ -57,14 +57,14 @@ async def test_list( if legacy_replica: gateway_replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, populate_configuration=populate_configuration, ) gateway.gateway_replica_id = gateway_replica.id # pre-0.20.25 relationship style else: gateway_replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -143,14 +143,14 @@ async def test_get( if legacy_replica: gateway_replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, populate_configuration=populate_configuration, ) gateway.gateway_replica_id = gateway_replica.id # pre-0.20.25 relationship style else: gateway_replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -220,7 +220,7 @@ async def test_list_legacy_client_populates_compat_fields( ) gateway_replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, ) response = await client.post( @@ -249,7 +249,7 @@ async def test_list_non_member_public_project( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, ) response = await client.post( @@ -275,7 +275,7 @@ async def test_get_non_member_public_project( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, ) response = await client.post( @@ -326,7 +326,7 @@ async def test_list_returns_imported_gateway_with_include_imported( backend_id=backend.id, name="exported-gateway", ) - await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend=backend, gateway_id=gateway.id) await create_export( session=session, exporter_project=exporter_project, @@ -369,7 +369,7 @@ async def test_list_not_returns_imported_gateway_without_include_imported( backend_id=backend.id, name="exported-gateway", ) - await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend=backend, gateway_id=gateway.id) await create_export( session=session, exporter_project=exporter_project, @@ -410,7 +410,7 @@ async def test_get_returns_imported_gateway( backend_id=backend.id, name="exported-gateway", ) - await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend=backend, gateway_id=gateway.id) await create_export( session=session, exporter_project=exporter_project, @@ -457,7 +457,7 @@ async def test_get_returns_403_on_foreign_gateway_if_not_imported( backend_id=backend.id, name="exported-gateway", ) - await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend=backend, gateway_id=gateway.id) await create_export( session=session, exporter_project=exporter_project, @@ -820,7 +820,7 @@ async def test_set_default_gateway( ) gateway_replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -889,7 +889,7 @@ async def test_set_default_gateway( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=second_gateway.id, populate_configuration=populate_configuration, ) @@ -994,7 +994,7 @@ async def test_set_imported_gateway_as_default( backend_id=backend.id, name="exported-gateway", ) - await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend=backend, gateway_id=gateway.id) await create_export( session=session, exporter_project=exporter_project, @@ -1038,7 +1038,7 @@ async def test_cannot_set_non_imported_foreign_gateway_as_default( backend_id=backend.id, name="exported-gateway", ) - await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend=backend, gateway_id=gateway.id) await create_export( session=session, exporter_project=exporter_project, @@ -1097,7 +1097,7 @@ async def test_marks_gateways_to_be_deleted( ) gateway_replica_aws = await create_gateway_replica( session=session, - backend_id=backend_aws.id, + backend=backend_aws, gateway_id=gateway_aws.id, populate_configuration=populate_configuration, ) @@ -1110,7 +1110,7 @@ async def test_marks_gateways_to_be_deleted( ) gateway_replica_gcp = await create_gateway_replica( session=session, - backend_id=backend_gcp.id, + backend=backend_gcp, gateway_id=gateway_gcp.id, populate_configuration=populate_configuration, ) @@ -1223,7 +1223,7 @@ async def test_set_wildcard_domain( ) gateway_replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -1468,7 +1468,7 @@ async def test_get_plan_with_existing_gateway_no_changes( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -1514,7 +1514,7 @@ async def test_get_plan_with_domain_change_is_update( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -1562,7 +1562,7 @@ async def test_get_plan_rejects_failed_gateway( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -1606,7 +1606,7 @@ async def test_get_plan_with_region_change_is_create( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -1681,7 +1681,7 @@ async def test_get_plan_rejects_to_be_deleted_gateway( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -1965,7 +1965,7 @@ async def test_updates_in_place( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -2027,7 +2027,7 @@ async def test_updates_in_place_with_force_apply( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -2082,7 +2082,7 @@ async def test_force_apply_no_changes_succeeds( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -2130,7 +2130,7 @@ async def test_rejects_update( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -2184,7 +2184,7 @@ async def test_rejects_update_with_force_apply( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -2231,7 +2231,7 @@ async def test_returns_error_on_missing_current_resource( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -2279,7 +2279,7 @@ async def test_returns_error_on_current_resource_mismatch( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -2336,7 +2336,7 @@ async def test_rejects_apply_on_to_be_deleted_gateway( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -2401,7 +2401,7 @@ async def test_rejects_apply_on_failed_gateway( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -2456,9 +2456,7 @@ async def test_sets_default_in_place( region="us-east-1", populate_configuration=populate_configuration, ) - await create_gateway_replica( - session=session, backend_id=backend.id, gateway_id=first_gateway.id - ) + await create_gateway_replica(session=session, backend=backend, gateway_id=first_gateway.id) second_gateway = await create_gateway( session=session, project_id=project.id, @@ -2468,7 +2466,7 @@ async def test_sets_default_in_place( populate_configuration=populate_configuration, ) await create_gateway_replica( - session=session, backend_id=backend.id, gateway_id=second_gateway.id + session=session, backend=backend, gateway_id=second_gateway.id ) response = await client.post( f"/api/project/{project.name}/gateways/set_default", @@ -2539,7 +2537,7 @@ async def test_unsets_default_in_place( region="us-east-1", populate_configuration=populate_configuration, ) - await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend=backend, gateway_id=gateway.id) response = await client.post( f"/api/project/{project.name}/gateways/set_default", json={"name": gateway.name}, @@ -2604,7 +2602,7 @@ async def test_omitted_default_leaves_current_status_unchanged( name="my-gateway", region="us-east-1", ) - await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend=backend, gateway_id=gateway.id) if initial_default: response = await client.post( f"/api/project/{project.name}/gateways/set_default", @@ -2669,7 +2667,7 @@ async def test_legacy_client_default_false_is_treated_as_omitted( name="my-gateway", region="us-east-1", ) - await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend=backend, gateway_id=gateway.id) if initial_default: response = await client.post( f"/api/project/{project.name}/gateways/set_default", diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py index 832633db9..b1e3eef18 100644 --- a/src/tests/_internal/server/routers/test_runs.py +++ b/src/tests/_internal/server/routers/test_runs.py @@ -3976,7 +3976,7 @@ async def test_submit_to_correct_proxy( ) await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, ) if is_default: @@ -4043,7 +4043,7 @@ async def test_submit_to_gateway_by_name( if legacy_replica: gateway_replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, populate_configuration=populate_configuration, ) gateway.gateway_replica_id = gateway_replica.id @@ -4051,7 +4051,7 @@ async def test_submit_to_gateway_by_name( else: gateway_replica = await create_gateway_replica( session=session, - backend_id=backend.id, + backend=backend, gateway_id=gateway.id, populate_configuration=populate_configuration, ) @@ -4162,7 +4162,7 @@ async def test_submit_to_foreign_gateway_only_if_imported( name="exported-gateway", wildcard_domain="exported-gateway.example", ) - await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend=backend, gateway_id=gateway.id) importer_user = await create_user( session=session, global_role=GlobalRole.USER, name="importer_user" @@ -4256,7 +4256,7 @@ async def test_not_submits_to_default_gateway_if_not_imported( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend=backend, gateway_id=gateway.id) service_project = await create_project(session=session, owner=user, name="service-project") # The project's default_gateway_id may point to the gateway (e.g., if the gateway was @@ -4316,7 +4316,7 @@ async def test_interpolates_project_name_in_imported_gateway_domain( name="exported-gateway", wildcard_domain="${{ run.project_name }}.example.com", ) - await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend=backend, gateway_id=gateway.id) importer_user = await create_user( session=session, global_role=GlobalRole.USER, name="importer_user" @@ -4380,7 +4380,7 @@ async def test_returns_error_if_imported_gateway_domain_has_unknown_variable( name="exported-gateway", wildcard_domain="${{ run.unknown_variable }}.example.com", ) - await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend=backend, gateway_id=gateway.id) importer_user = await create_user( session=session, global_role=GlobalRole.USER, name="importer_user" @@ -4451,7 +4451,7 @@ async def test_return_error_if_default_gateway_forbids_new_services( wildcard_domain="example.com", forbid_new_services=True, ) - await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend=backend, gateway_id=gateway.id) project.default_gateway_id = gateway.id await session.commit() @@ -4495,7 +4495,7 @@ async def test_return_error_if_explicitly_specified_gateway_forbids_new_services wildcard_domain="example.com", forbid_new_services=True, ) - await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend=backend, gateway_id=gateway.id) response = await client.post( f"/api/project/{project.name}/runs/apply", diff --git a/src/tests/_internal/server/services/gateways/test_gateways.py b/src/tests/_internal/server/services/gateways/test_gateways.py index 0e4fbe938..816619a4e 100644 --- a/src/tests/_internal/server/services/gateways/test_gateways.py +++ b/src/tests/_internal/server/services/gateways/test_gateways.py @@ -56,7 +56,7 @@ async def test_new_style_returns_gateway_replicas(self, test_db, session: AsyncS session=session, project_id=project.id, backend_id=backend.id ) replica = await create_gateway_replica( - session=session, gateway_id=gateway.id, backend_id=backend.id + session=session, gateway_id=gateway.id, backend=backend ) await session.refresh(gateway, ["gateway_replicas", "gateway_replica"]) result = get_gateway_replica_models(gateway) @@ -66,7 +66,7 @@ async def test_new_style_returns_gateway_replicas(self, test_db, session: AsyncS async def test_old_style_returns_single_replica(self, test_db, session: AsyncSession): project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) - replica = await create_gateway_replica(session=session, backend_id=backend.id) + replica = await create_gateway_replica(session=session, backend=backend) gateway = await create_gateway( session=session, project_id=project.id, backend_id=backend.id )