diff --git a/changes/11643.fix.md b/changes/11643.fix.md new file mode 100644 index 00000000000..8c19d04a2b4 --- /dev/null +++ b/changes/11643.fix.md @@ -0,0 +1 @@ +Enforce `ON DELETE RESTRICT` on `images.registry_id` so deleting a container registry that still has images is blocked at the database level instead of leaving dangling references that surfaced later as a misleading "Image not found in database" error. diff --git a/src/ai/backend/manager/models/alembic/versions/045b0e0e4384_add_fk_from_images_registry_id_to_.py b/src/ai/backend/manager/models/alembic/versions/045b0e0e4384_add_fk_from_images_registry_id_to_.py new file mode 100644 index 00000000000..61d0faebb2e --- /dev/null +++ b/src/ai/backend/manager/models/alembic/versions/045b0e0e4384_add_fk_from_images_registry_id_to_.py @@ -0,0 +1,55 @@ +"""add FK to images.registry_id (RESTRICT) + +Adds a foreign key from ``images.registry_id`` to ``container_registries.id`` +with ``ON DELETE RESTRICT``. + +Rationale: the column previously carried a UUID with no DB-level FK, so +deleting a container registry left dangling references on the images +table. The sokovan scheduler would later surface this as a misleading +"Image not found in database" error during session launch. RESTRICT +prevents the registry deletion at the database level, forcing callers +to ``clear_images`` (soft-delete) or ``admin_purge`` the dependent rows +first. CASCADE was rejected because it would also hard-delete image +rows referenced by historical kernels / deployment_revisions via +SET NULL, erasing audit traceability. + +Pre-existing dangling rows (image rows whose ``registry_id`` no longer +resolves) are hard-deleted before the constraint is created. The column +remains ``NOT NULL``, so a SET NULL backfill is not an option, and a +sentinel re-assignment would mask the original deletion intent. + +Revision ID: 045b0e0e4384 +Revises: ba42cb865efe +Create Date: 2026-05-17 + +""" + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "045b0e0e4384" +down_revision = "ba42cb865efe" +# Part of: 26.5.0 +branch_labels = None +depends_on = None + +_FK_NAME = "fk_images_registry_id_container_registries" + + +def upgrade() -> None: + op.execute( + sa.text("DELETE FROM images WHERE registry_id NOT IN (SELECT id FROM container_registries)") + ) + op.create_foreign_key( + _FK_NAME, + source_table="images", + referent_table="container_registries", + local_cols=["registry_id"], + remote_cols=["id"], + ondelete="RESTRICT", + ) + + +def downgrade() -> None: + op.drop_constraint(_FK_NAME, "images", type_="foreignkey") diff --git a/src/ai/backend/manager/models/image/row.py b/src/ai/backend/manager/models/image/row.py index c014d4c80aa..b453b81299f 100644 --- a/src/ai/backend/manager/models/image/row.py +++ b/src/ai/backend/manager/models/image/row.py @@ -334,7 +334,13 @@ class ImageRow(Base): # type: ignore[misc] ) tag: Mapped[str | None] = mapped_column("tag", sa.TEXT, nullable=True) registry: Mapped[str] = mapped_column("registry", sa.String, nullable=False, index=True) - registry_id: Mapped[UUID] = mapped_column("registry_id", GUID, nullable=False, index=True) + registry_id: Mapped[UUID] = mapped_column( + "registry_id", + GUID, + sa.ForeignKey("container_registries.id", ondelete="RESTRICT"), + nullable=False, + index=True, + ) architecture: Mapped[str] = mapped_column( "architecture", sa.String, nullable=False, index=True, default="x86_64" ) diff --git a/tests/unit/manager/models/domain/test_conditions.py b/tests/unit/manager/models/domain/test_conditions.py index 2a523af611f..b0efb7d4796 100644 --- a/tests/unit/manager/models/domain/test_conditions.py +++ b/tests/unit/manager/models/domain/test_conditions.py @@ -15,6 +15,7 @@ from ai.backend.manager.data.group.types import ProjectType from ai.backend.manager.data.user.types import UserStatus from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import ( DeploymentAutoScalingPolicyRow, ) @@ -70,6 +71,7 @@ def _make_password_info() -> PasswordInfo: UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/models/group/test_conditions.py b/tests/unit/manager/models/group/test_conditions.py index f4080afbf68..125b2697217 100644 --- a/tests/unit/manager/models/group/test_conditions.py +++ b/tests/unit/manager/models/group/test_conditions.py @@ -16,6 +16,7 @@ from ai.backend.manager.data.group.types import ProjectType from ai.backend.manager.data.user.types import UserStatus from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import ( DeploymentAutoScalingPolicyRow, ) @@ -62,6 +63,7 @@ KeyPairRow, GroupRow, AssocGroupUserRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/models/scaling_group/test_conditions.py b/tests/unit/manager/models/scaling_group/test_conditions.py index 251cb4d7e3a..054b56dbd69 100644 --- a/tests/unit/manager/models/scaling_group/test_conditions.py +++ b/tests/unit/manager/models/scaling_group/test_conditions.py @@ -8,6 +8,7 @@ import sqlalchemy as sa from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -47,6 +48,7 @@ UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/models/test_deployment_auto_scaling_policy.py b/tests/unit/manager/models/test_deployment_auto_scaling_policy.py index 52af097eb4d..10e4d2c4a2f 100644 --- a/tests/unit/manager/models/test_deployment_auto_scaling_policy.py +++ b/tests/unit/manager/models/test_deployment_auto_scaling_policy.py @@ -10,6 +10,7 @@ import pytest import sqlalchemy as sa +from ai.backend.common.container_registry import ContainerRegistryType from ai.backend.common.data.endpoint.types import EndpointLifecycle from ai.backend.common.types import ( AutoScalingMetricComparator, @@ -20,6 +21,7 @@ from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm from ai.backend.manager.data.image.types import ImageType from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import ( DeploymentAutoScalingPolicyData, DeploymentAutoScalingPolicyRow, @@ -90,6 +92,7 @@ async def db_with_cleanup( KeyPairRow, GroupRow, VFolderRow, + ContainerRegistryRow, ImageRow, SessionRow, KernelRow, @@ -224,13 +227,23 @@ async def test_image( db_with_cleanup: ExtendedAsyncSAEngine, ) -> AsyncGenerator[ImageRow, None]: """Create test image.""" + registry_id = uuid.uuid4() async with db_with_cleanup.begin_session() as db_sess: + db_sess.add( + ContainerRegistryRow( + id=registry_id, + url="https://docker.io", + registry_name=f"reg-{uuid.uuid4().hex[:8]}", + type=ContainerRegistryType.DOCKER, + ) + ) + await db_sess.flush() image = ImageRow( name="test-image:latest", project=str(uuid.uuid4()), image="test-image", registry="docker.io", - registry_id=uuid.uuid4(), + registry_id=registry_id, architecture="x86_64", is_local=False, config_digest="sha256:abc123", diff --git a/tests/unit/manager/models/test_deployment_revision.py b/tests/unit/manager/models/test_deployment_revision.py index cb6b123879d..7690f74c91b 100644 --- a/tests/unit/manager/models/test_deployment_revision.py +++ b/tests/unit/manager/models/test_deployment_revision.py @@ -11,6 +11,7 @@ import sqlalchemy as sa from ai.backend.common.config import ModelDefinitionDraft +from ai.backend.common.container_registry import ContainerRegistryType from ai.backend.common.data.endpoint.types import EndpointLifecycle from ai.backend.common.identifier.vfolder import VFolderUUID from ai.backend.common.types import ( @@ -25,6 +26,7 @@ from ai.backend.manager.data.deployment.types import ModelRevisionData from ai.backend.manager.data.image.types import ImageType from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -95,6 +97,7 @@ async def db_with_cleanup( GroupRow, AgentRow, VFolderRow, + ContainerRegistryRow, ImageRow, ResourcePresetRow, ResourceSlotTypeRow, @@ -242,13 +245,23 @@ async def test_image( db_with_cleanup: ExtendedAsyncSAEngine, ) -> AsyncGenerator[ImageRow, None]: """Create test image.""" + registry_id = uuid.uuid4() async with db_with_cleanup.begin_session() as db_sess: + db_sess.add( + ContainerRegistryRow( + id=registry_id, + url="https://docker.io", + registry_name=f"reg-{uuid.uuid4().hex[:8]}", + type=ContainerRegistryType.DOCKER, + ) + ) + await db_sess.flush() image = ImageRow( name="test-image:latest", project=str(uuid.uuid4()), image="test-image", registry="docker.io", - registry_id=uuid.uuid4(), + registry_id=registry_id, architecture="x86_64", is_local=False, config_digest="sha256:abc123", diff --git a/tests/unit/manager/models/test_session.py b/tests/unit/manager/models/test_session.py index ec98779c239..56aad948911 100644 --- a/tests/unit/manager/models/test_session.py +++ b/tests/unit/manager/models/test_session.py @@ -10,6 +10,7 @@ from ai.backend.common.types import BinarySize, ResourceSlot from ai.backend.manager.data.session.types import SessionStatus from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -80,6 +81,7 @@ async def database_with_tables( GroupRow, AgentRow, VFolderRow, + ContainerRegistryRow, ImageRow, ResourcePresetRow, EndpointRow, diff --git a/tests/unit/manager/models/test_vfolder_prepare_mounts.py b/tests/unit/manager/models/test_vfolder_prepare_mounts.py index 1a0276ab718..c13399cc871 100644 --- a/tests/unit/manager/models/test_vfolder_prepare_mounts.py +++ b/tests/unit/manager/models/test_vfolder_prepare_mounts.py @@ -20,6 +20,7 @@ from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm from ai.backend.manager.errors.api import InvalidAPIParameters from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -128,6 +129,7 @@ async def db_with_cleanup( AgentRow, VFolderRow, VFolderPermissionRow, + ContainerRegistryRow, ImageRow, ResourcePresetRow, RuntimeVariantRow, diff --git a/tests/unit/manager/models/test_vfolder_quota_scope.py b/tests/unit/manager/models/test_vfolder_quota_scope.py index e7c09c251c6..1830e3e1202 100644 --- a/tests/unit/manager/models/test_vfolder_quota_scope.py +++ b/tests/unit/manager/models/test_vfolder_quota_scope.py @@ -13,6 +13,7 @@ ) from ai.backend.manager.errors.api import InvalidAPIParameters from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -74,6 +75,7 @@ async def db_with_cleanup( GroupRow, AgentRow, VFolderRow, + ContainerRegistryRow, ImageRow, ResourcePresetRow, EndpointRow, diff --git a/tests/unit/manager/models/user/test_conditions.py b/tests/unit/manager/models/user/test_conditions.py index a44f5a56f29..05729b2357e 100644 --- a/tests/unit/manager/models/user/test_conditions.py +++ b/tests/unit/manager/models/user/test_conditions.py @@ -14,6 +14,7 @@ from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm from ai.backend.manager.data.group.types import ProjectType from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import ( DeploymentAutoScalingPolicyRow, ) @@ -60,6 +61,7 @@ KeyPairRow, GroupRow, AssocGroupUserRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/agent/test_repository.py b/tests/unit/manager/repositories/agent/test_repository.py index d8929e96ea8..1ac01c2ee52 100644 --- a/tests/unit/manager/repositories/agent/test_repository.py +++ b/tests/unit/manager/repositories/agent/test_repository.py @@ -40,6 +40,7 @@ from ai.backend.manager.data.session.types import SessionStatus from ai.backend.manager.errors.resource import ScalingGroupNotFound from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -110,6 +111,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, @@ -620,6 +622,7 @@ async def db_with_tables( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/app_config/test_app_config.py b/tests/unit/manager/repositories/app_config/test_app_config.py index 9823ecc6def..8aa9c7d3fa5 100644 --- a/tests/unit/manager/repositories/app_config/test_app_config.py +++ b/tests/unit/manager/repositories/app_config/test_app_config.py @@ -15,6 +15,7 @@ from ai.backend.manager.data.permission.types import RBACElementRef from ai.backend.manager.models.agent import AgentRow from ai.backend.manager.models.app_config import AppConfigRow, AppConfigScopeType +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -84,6 +85,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/artifact/test_artifact_repository.py b/tests/unit/manager/repositories/artifact/test_artifact_repository.py index c9ec15c6cf4..84bfb341bce 100644 --- a/tests/unit/manager/repositories/artifact/test_artifact_repository.py +++ b/tests/unit/manager/repositories/artifact/test_artifact_repository.py @@ -23,6 +23,7 @@ from ai.backend.manager.models.agent import AgentRow from ai.backend.manager.models.artifact import ArtifactRow from ai.backend.manager.models.artifact_revision import ArtifactRevisionRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -82,6 +83,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/artifact_revision/test_artifact_revision_repository.py b/tests/unit/manager/repositories/artifact_revision/test_artifact_revision_repository.py index 7f5e2782cf7..3410dd39ac3 100644 --- a/tests/unit/manager/repositories/artifact_revision/test_artifact_revision_repository.py +++ b/tests/unit/manager/repositories/artifact_revision/test_artifact_revision_repository.py @@ -24,6 +24,7 @@ from ai.backend.manager.models.agent import AgentRow from ai.backend.manager.models.artifact import ArtifactRow from ai.backend.manager.models.artifact_revision import ArtifactRevisionRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -87,6 +88,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/auth/test_auth_repository.py b/tests/unit/manager/repositories/auth/test_auth_repository.py index 23937f65f57..6992c8f5faa 100644 --- a/tests/unit/manager/repositories/auth/test_auth_repository.py +++ b/tests/unit/manager/repositories/auth/test_auth_repository.py @@ -22,6 +22,7 @@ from ai.backend.manager.data.permission.types import EntityType, ScopeType from ai.backend.manager.errors.auth import GroupMembershipNotFoundError from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -100,6 +101,7 @@ async def db_with_cleanup( GroupRow, AssocGroupUserRow, AssociationScopesEntitiesRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/domain/test_domain_repository.py b/tests/unit/manager/repositories/domain/test_domain_repository.py index bc1f87a66c3..ef4cc080b55 100644 --- a/tests/unit/manager/repositories/domain/test_domain_repository.py +++ b/tests/unit/manager/repositories/domain/test_domain_repository.py @@ -27,6 +27,7 @@ DomainHasUsers, ) from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -82,6 +83,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/error_log/test_error_log_repository.py b/tests/unit/manager/repositories/error_log/test_error_log_repository.py index 59817ef1d94..bf1f1f849d0 100644 --- a/tests/unit/manager/repositories/error_log/test_error_log_repository.py +++ b/tests/unit/manager/repositories/error_log/test_error_log_repository.py @@ -16,6 +16,7 @@ ErrorLogSeverity, ) from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -80,6 +81,7 @@ async def db_with_cleanup( GroupRow, AgentRow, VFolderRow, + ContainerRegistryRow, ImageRow, ResourcePresetRow, EndpointRow, diff --git a/tests/unit/manager/repositories/fair_share/test_fair_share_repository.py b/tests/unit/manager/repositories/fair_share/test_fair_share_repository.py index d834aacee98..ee88ccff9dd 100644 --- a/tests/unit/manager/repositories/fair_share/test_fair_share_repository.py +++ b/tests/unit/manager/repositories/fair_share/test_fair_share_repository.py @@ -14,6 +14,7 @@ from ai.backend.common.types import ResourceSlot from ai.backend.manager.errors.resource import DomainNotFound from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.domain import DomainRow from ai.backend.manager.models.fair_share import ( DomainFairShareRow, @@ -91,6 +92,7 @@ async def db_with_cleanup( ScalingGroupForProjectRow, AssocGroupUserRow, AgentRow, + ContainerRegistryRow, ImageRow, SessionRow, KernelRow, diff --git a/tests/unit/manager/repositories/group/test_assign_users_to_project.py b/tests/unit/manager/repositories/group/test_assign_users_to_project.py index 74d69ee2c8a..e0045383f8d 100644 --- a/tests/unit/manager/repositories/group/test_assign_users_to_project.py +++ b/tests/unit/manager/repositories/group/test_assign_users_to_project.py @@ -13,6 +13,7 @@ from ai.backend.manager.data.group.types import ProjectType from ai.backend.manager.data.permission.types import EntityType, ScopeType from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.domain import DomainRow from ai.backend.manager.models.endpoint import EndpointRow from ai.backend.manager.models.group import GroupRow @@ -73,6 +74,7 @@ async def db_with_cleanup( KeyPairRow, GroupRow, AssociationScopesEntitiesRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, @@ -462,6 +464,7 @@ async def db_with_cleanup( KeyPairRow, GroupRow, AssociationScopesEntitiesRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/group/test_group_db_source.py b/tests/unit/manager/repositories/group/test_group_db_source.py index db0b6ea82d8..65b767a33fb 100644 --- a/tests/unit/manager/repositories/group/test_group_db_source.py +++ b/tests/unit/manager/repositories/group/test_group_db_source.py @@ -16,6 +16,7 @@ from ai.backend.manager.data.model_serving.types import EndpointLifecycle from ai.backend.manager.errors.resource import ProjectHasActiveEndpointsError from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -90,6 +91,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/group/test_group_repository.py b/tests/unit/manager/repositories/group/test_group_repository.py index 1d03f899c1f..a28fa71e042 100644 --- a/tests/unit/manager/repositories/group/test_group_repository.py +++ b/tests/unit/manager/repositories/group/test_group_repository.py @@ -39,6 +39,7 @@ ProjectNotFound, ) from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -258,6 +259,7 @@ async def db_with_cleanup( KeyPairRow, GroupRow, AssociationScopesEntitiesRow, # RBAC scopes-entities association + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/group/test_unassign_users_from_project.py b/tests/unit/manager/repositories/group/test_unassign_users_from_project.py index 33e43a1dd81..5425872eee1 100644 --- a/tests/unit/manager/repositories/group/test_unassign_users_from_project.py +++ b/tests/unit/manager/repositories/group/test_unassign_users_from_project.py @@ -13,6 +13,7 @@ from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm from ai.backend.manager.data.group.types import ProjectType from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.domain import DomainRow from ai.backend.manager.models.endpoint import EndpointRow from ai.backend.manager.models.group import GroupRow @@ -72,6 +73,7 @@ async def db_with_cleanup( KeyPairRow, GroupRow, AssociationScopesEntitiesRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/keypair_resource_policy/test_keypair_resource_policy.py b/tests/unit/manager/repositories/keypair_resource_policy/test_keypair_resource_policy.py index 1cce4bc8ed4..df3492e3f98 100644 --- a/tests/unit/manager/repositories/keypair_resource_policy/test_keypair_resource_policy.py +++ b/tests/unit/manager/repositories/keypair_resource_policy/test_keypair_resource_policy.py @@ -14,6 +14,7 @@ VFolderHostPermission, ) from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -76,6 +77,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/model_card/test_model_card_creator.py b/tests/unit/manager/repositories/model_card/test_model_card_creator.py index 7ce38a6e26b..bfa7d6499dd 100644 --- a/tests/unit/manager/repositories/model_card/test_model_card_creator.py +++ b/tests/unit/manager/repositories/model_card/test_model_card_creator.py @@ -13,6 +13,7 @@ from ai.backend.manager.data.model_card.types import ResourceRequirementEntry from ai.backend.manager.data.permission.types import RBACElementRef, RBACElementType from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.domain import DomainRow from ai.backend.manager.models.group import GroupRow from ai.backend.manager.models.hasher.types import PasswordInfo @@ -72,6 +73,7 @@ async def db_with_cleanup( KeyPairRow, GroupRow, AgentRow, + ContainerRegistryRow, ImageRow, VFolderRow, SessionRow, diff --git a/tests/unit/manager/repositories/model_card/test_model_card_delete.py b/tests/unit/manager/repositories/model_card/test_model_card_delete.py index 84e88d514dc..483e495479e 100644 --- a/tests/unit/manager/repositories/model_card/test_model_card_delete.py +++ b/tests/unit/manager/repositories/model_card/test_model_card_delete.py @@ -28,6 +28,7 @@ from ai.backend.manager.data.session.types import SessionStatus from ai.backend.manager.data.vfolder.types import VFolderOperationStatus from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.domain import DomainRow from ai.backend.manager.models.group import GroupRow from ai.backend.manager.models.hasher.types import PasswordInfo @@ -123,6 +124,7 @@ async def db_with_cleanup( KeyPairRow, GroupRow, AgentRow, + ContainerRegistryRow, ImageRow, VFolderRow, SessionRow, diff --git a/tests/unit/manager/repositories/model_card/test_model_card_scan.py b/tests/unit/manager/repositories/model_card/test_model_card_scan.py index 09470fd582d..19cb83a6d4a 100644 --- a/tests/unit/manager/repositories/model_card/test_model_card_scan.py +++ b/tests/unit/manager/repositories/model_card/test_model_card_scan.py @@ -26,6 +26,7 @@ from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm from ai.backend.manager.data.model_card.types import ResourceRequirementEntry from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_revision_preset.row import DeploymentRevisionPresetRow from ai.backend.manager.models.domain import DomainRow from ai.backend.manager.models.group import GroupRow @@ -89,6 +90,7 @@ async def db_with_cleanup( KeyPairRow, GroupRow, AgentRow, + ContainerRegistryRow, ImageRow, VFolderRow, SessionRow, diff --git a/tests/unit/manager/repositories/model_card/test_model_card_search_in_project.py b/tests/unit/manager/repositories/model_card/test_model_card_search_in_project.py index 7c20a1b8398..1deea2de2e0 100644 --- a/tests/unit/manager/repositories/model_card/test_model_card_search_in_project.py +++ b/tests/unit/manager/repositories/model_card/test_model_card_search_in_project.py @@ -19,6 +19,7 @@ from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm from ai.backend.manager.errors.common import GenericForbidden from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.domain import DomainRow from ai.backend.manager.models.group import GroupRow from ai.backend.manager.models.hasher.types import PasswordInfo @@ -71,6 +72,7 @@ async def db_with_cleanup( KeyPairRow, GroupRow, AgentRow, + ContainerRegistryRow, ImageRow, VFolderRow, SessionRow, diff --git a/tests/unit/manager/repositories/notification/test_notification_options.py b/tests/unit/manager/repositories/notification/test_notification_options.py index 5341de55f39..093592abde5 100644 --- a/tests/unit/manager/repositories/notification/test_notification_options.py +++ b/tests/unit/manager/repositories/notification/test_notification_options.py @@ -20,6 +20,7 @@ ) from ai.backend.common.types import BinarySize, ResourceSlot from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -97,6 +98,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, @@ -984,6 +986,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/notification/test_notification_repository.py b/tests/unit/manager/repositories/notification/test_notification_repository.py index 74004c52ea7..4fcf61678a3 100644 --- a/tests/unit/manager/repositories/notification/test_notification_repository.py +++ b/tests/unit/manager/repositories/notification/test_notification_repository.py @@ -24,6 +24,7 @@ NotificationRuleNotFound, ) from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -105,6 +106,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/permission_controller/test_role_assignment.py b/tests/unit/manager/repositories/permission_controller/test_role_assignment.py index 0f8b8ed67e1..6c142c48c50 100644 --- a/tests/unit/manager/repositories/permission_controller/test_role_assignment.py +++ b/tests/unit/manager/repositories/permission_controller/test_role_assignment.py @@ -17,6 +17,7 @@ ) from ai.backend.manager.data.permission.types import EntityType, ScopeType from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.domain import DomainRow from ai.backend.manager.models.endpoint import EndpointRow from ai.backend.manager.models.group import GroupRow @@ -79,6 +80,7 @@ async def db_with_cleanup( KeyPairRow, GroupRow, AssociationScopesEntitiesRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/permission_controller/test_role_invitation.py b/tests/unit/manager/repositories/permission_controller/test_role_invitation.py index 42878095a3f..3598252ae42 100644 --- a/tests/unit/manager/repositories/permission_controller/test_role_invitation.py +++ b/tests/unit/manager/repositories/permission_controller/test_role_invitation.py @@ -18,6 +18,7 @@ RoleInvitationNotFound, ) from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.domain import DomainRow from ai.backend.manager.models.endpoint import EndpointRow from ai.backend.manager.models.group import GroupRow @@ -80,6 +81,7 @@ async def db( GroupRow, AssocGroupUserRow, AssociationScopesEntitiesRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/permission_controller/test_search_scopes.py b/tests/unit/manager/repositories/permission_controller/test_search_scopes.py index 16dbf5513e9..055d44615a9 100644 --- a/tests/unit/manager/repositories/permission_controller/test_search_scopes.py +++ b/tests/unit/manager/repositories/permission_controller/test_search_scopes.py @@ -14,6 +14,7 @@ from ai.backend.common.data.permission.types import RBACElementType, ScopeType from ai.backend.common.types import ResourceSlot from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -88,6 +89,7 @@ async def db_with_scope_tables( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, @@ -342,6 +344,7 @@ async def db_with_scope_tables( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, @@ -556,6 +559,7 @@ async def db_with_scope_tables( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, @@ -770,6 +774,7 @@ async def db_with_scope_tables( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, @@ -819,6 +824,7 @@ async def db_with_scope_tables( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/permission_controller/test_user_scope_sync.py b/tests/unit/manager/repositories/permission_controller/test_user_scope_sync.py index 7b13413f15e..ef914771209 100644 --- a/tests/unit/manager/repositories/permission_controller/test_user_scope_sync.py +++ b/tests/unit/manager/repositories/permission_controller/test_user_scope_sync.py @@ -17,6 +17,7 @@ ) from ai.backend.manager.data.permission.types import EntityType, ScopeType from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.domain import DomainRow from ai.backend.manager.models.endpoint import EndpointRow from ai.backend.manager.models.group import GroupRow @@ -79,6 +80,7 @@ async def db_with_cleanup( GroupRow, AssocGroupUserRow, AssociationScopesEntitiesRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/project_resource_policy/test_project_resource_policy.py b/tests/unit/manager/repositories/project_resource_policy/test_project_resource_policy.py index 9b38347d169..6ff0e66f297 100644 --- a/tests/unit/manager/repositories/project_resource_policy/test_project_resource_policy.py +++ b/tests/unit/manager/repositories/project_resource_policy/test_project_resource_policy.py @@ -17,6 +17,7 @@ from ai.backend.manager.errors.common import ObjectNotFound from ai.backend.manager.errors.repository import RepositoryIntegrityError from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -83,6 +84,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/resource_preset/test_check_presets.py b/tests/unit/manager/repositories/resource_preset/test_check_presets.py index 75e8b8eaf31..36cc9d1feca 100644 --- a/tests/unit/manager/repositories/resource_preset/test_check_presets.py +++ b/tests/unit/manager/repositories/resource_preset/test_check_presets.py @@ -39,6 +39,7 @@ from ai.backend.manager.data.kernel.types import KernelStatus from ai.backend.manager.data.session.types import SessionStatus from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -118,6 +119,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, @@ -1196,6 +1198,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/resource_preset/test_resource_preset_cache_invalidation.py b/tests/unit/manager/repositories/resource_preset/test_resource_preset_cache_invalidation.py index 0d202f8d27f..5aa9f124a4a 100644 --- a/tests/unit/manager/repositories/resource_preset/test_resource_preset_cache_invalidation.py +++ b/tests/unit/manager/repositories/resource_preset/test_resource_preset_cache_invalidation.py @@ -16,6 +16,7 @@ from ai.backend.common.types import AccessKey, BinarySize, ResourceSlot, ValkeyTarget from ai.backend.manager.data.resource_preset.types import ResourcePresetData from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -75,6 +76,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/resource_usage_history/test_resource_usage_history_repository.py b/tests/unit/manager/repositories/resource_usage_history/test_resource_usage_history_repository.py index 88505bc5073..12130ea9968 100644 --- a/tests/unit/manager/repositories/resource_usage_history/test_resource_usage_history_repository.py +++ b/tests/unit/manager/repositories/resource_usage_history/test_resource_usage_history_repository.py @@ -14,6 +14,7 @@ from ai.backend.common.types import ResourceSlot from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.domain import DomainRow from ai.backend.manager.models.group import GroupRow from ai.backend.manager.models.image import ImageRow @@ -85,6 +86,7 @@ async def db_with_cleanup( KeyPairRow, GroupRow, AgentRow, + ContainerRegistryRow, ImageRow, SessionRow, KernelRow, diff --git a/tests/unit/manager/repositories/resource_usage_history/test_usage_bucket_entries.py b/tests/unit/manager/repositories/resource_usage_history/test_usage_bucket_entries.py index 5287fbb4771..7f9cd5f3a14 100644 --- a/tests/unit/manager/repositories/resource_usage_history/test_usage_bucket_entries.py +++ b/tests/unit/manager/repositories/resource_usage_history/test_usage_bucket_entries.py @@ -16,6 +16,7 @@ from ai.backend.common.types import ResourceSlot from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.domain import DomainRow from ai.backend.manager.models.group import GroupRow from ai.backend.manager.models.image import ImageRow @@ -73,6 +74,7 @@ async def db_with_cleanup( KeyPairRow, GroupRow, AgentRow, + ContainerRegistryRow, ImageRow, SessionRow, KernelRow, diff --git a/tests/unit/manager/repositories/scaling_group/test_scaling_group_repository.py b/tests/unit/manager/repositories/scaling_group/test_scaling_group_repository.py index a154b9a7499..39c5d633648 100644 --- a/tests/unit/manager/repositories/scaling_group/test_scaling_group_repository.py +++ b/tests/unit/manager/repositories/scaling_group/test_scaling_group_repository.py @@ -16,6 +16,7 @@ from ai.backend.manager.defs import DEFAULT_ROLE from ai.backend.manager.errors.resource import ScalingGroupNotFound from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -108,6 +109,7 @@ async def db_with_cleanup( KeyPairRow, ScalingGroupForKeypairsRow, # depends on ScalingGroupRow and KeyPairRow GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/scheduler/test_session_network_propagation.py b/tests/unit/manager/repositories/scheduler/test_session_network_propagation.py index d2e048cfc95..be4c3888c71 100644 --- a/tests/unit/manager/repositories/scheduler/test_session_network_propagation.py +++ b/tests/unit/manager/repositories/scheduler/test_session_network_propagation.py @@ -26,6 +26,7 @@ from ai.backend.manager.data.session.types import SessionStatus from ai.backend.manager.data.user.types import UserStatus from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.domain import DomainRow from ai.backend.manager.models.group import GroupRow from ai.backend.manager.models.image.row import ImageRow @@ -64,6 +65,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, AgentRow, SessionRow, diff --git a/tests/unit/manager/repositories/scheduler/test_termination.py b/tests/unit/manager/repositories/scheduler/test_termination.py index 5ab47817fd6..64d6c87cd78 100644 --- a/tests/unit/manager/repositories/scheduler/test_termination.py +++ b/tests/unit/manager/repositories/scheduler/test_termination.py @@ -28,6 +28,7 @@ from ai.backend.manager.data.session.types import SessionStatus from ai.backend.manager.data.user.types import UserStatus from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy.row import ( DeploymentAutoScalingPolicyRow, ) @@ -79,6 +80,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/scheduling_history/test_scheduling_history_repository.py b/tests/unit/manager/repositories/scheduling_history/test_scheduling_history_repository.py index d9be960bcce..2f1693fd9e9 100644 --- a/tests/unit/manager/repositories/scheduling_history/test_scheduling_history_repository.py +++ b/tests/unit/manager/repositories/scheduling_history/test_scheduling_history_repository.py @@ -22,6 +22,7 @@ SessionStatus, ) from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import ( DeploymentAutoScalingPolicyRow, ) @@ -96,6 +97,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/test_utils.py b/tests/unit/manager/repositories/test_utils.py index 1b07d3d7abc..bb93f274b1a 100644 --- a/tests/unit/manager/repositories/test_utils.py +++ b/tests/unit/manager/repositories/test_utils.py @@ -14,6 +14,7 @@ from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm from ai.backend.manager.defs import DEFAULT_ROLE from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -76,6 +77,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/user/test_user_repository.py b/tests/unit/manager/repositories/user/test_user_repository.py index e06c9e2cc9e..d6c7eeedc5a 100644 --- a/tests/unit/manager/repositories/user/test_user_repository.py +++ b/tests/unit/manager/repositories/user/test_user_repository.py @@ -18,6 +18,7 @@ from ai.backend.manager.data.user.types import UserData from ai.backend.manager.errors.user import UserConflict, UserCreationBadRequest, UserNotFound from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -101,6 +102,7 @@ async def db_with_cleanup( KeyPairRow, GroupRow, AssocGroupUserRow, # Association table for users-groups + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/user_resource_policy/test_user_resource_policy_repository.py b/tests/unit/manager/repositories/user_resource_policy/test_user_resource_policy_repository.py index b302b93a651..cea292dd626 100644 --- a/tests/unit/manager/repositories/user_resource_policy/test_user_resource_policy_repository.py +++ b/tests/unit/manager/repositories/user_resource_policy/test_user_resource_policy_repository.py @@ -7,6 +7,7 @@ from ai.backend.common.exception import UserResourcePolicyNotFound from ai.backend.manager.data.resource.types import UserResourcePolicyData from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -69,6 +70,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, diff --git a/tests/unit/manager/repositories/vfolder/test_vfolder_repository.py b/tests/unit/manager/repositories/vfolder/test_vfolder_repository.py index 629ebeab2ff..5a1fece803b 100644 --- a/tests/unit/manager/repositories/vfolder/test_vfolder_repository.py +++ b/tests/unit/manager/repositories/vfolder/test_vfolder_repository.py @@ -35,6 +35,7 @@ ) from ai.backend.manager.errors.user import UserNotFound from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -139,6 +140,7 @@ async def db_with_cleanup( UserRow, KeyPairRow, GroupRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow, @@ -1348,6 +1350,7 @@ async def db_with_cleanup( KeyPairRow, GroupRow, VFolderRow, + ContainerRegistryRow, ImageRow, SessionRow, AgentRow, diff --git a/tests/unit/manager/test_query_userinfo.py b/tests/unit/manager/test_query_userinfo.py index 938e8ef8f59..c6a35173e5b 100644 --- a/tests/unit/manager/test_query_userinfo.py +++ b/tests/unit/manager/test_query_userinfo.py @@ -20,6 +20,7 @@ from ai.backend.manager.errors.api import InvalidAPIParameters from ai.backend.manager.errors.auth import AccessKeyNotFound from ai.backend.manager.models.agent import AgentRow +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow from ai.backend.manager.models.deployment_revision import DeploymentRevisionRow @@ -62,6 +63,7 @@ KeyPairRow, GroupRow, AssociationScopesEntitiesRow, + ContainerRegistryRow, ImageRow, VFolderRow, EndpointRow,