fix: update KeyPairData dotfiles type to bytes and adjust related logic in ArtifactDBSource and tests

Author: HyeockJinKim

src/ai/backend/manager/data/keypair/types.py

@@ -37,5 +37,5 @@ class KeyPairData:
     rate_limit: int
     ssh_public_key: Optional[str]
     ssh_private_key: Optional[str]
-    dotfiles: str
+    dotfiles: bytes
     bootstrap_script: str

src/ai/backend/manager/models/keypair/row.py

@@ -175,7 +175,7 @@ def to_data(self) -> KeyPairData:
             rate_limit=self.rate_limit if self.rate_limit is not None else 0,
             ssh_public_key=self.ssh_public_key,
             ssh_private_key=self.ssh_private_key,
-            dotfiles=self.dotfiles.decode("utf-8") if self.dotfiles else "",
+            dotfiles=self.dotfiles if self.dotfiles else b"\x90",
             bootstrap_script=self.bootstrap_script,
         )
 

src/ai/backend/manager/repositories/artifact/db_source/db_source.py

@@ -647,11 +647,7 @@ async def approve_artifact(self, revision_id: uuid.UUID) -> ArtifactRevisionData
             )
 
             result = await db_sess.execute(update_stmt)
-            updated_id = result.scalar_one_or_none()
-            if updated_id is None:
-                raise ArtifactUpdateError()
-
-            updated_row = await db_sess.get(ArtifactRevisionRow, updated_id)
+            updated_row = result.scalars().one_or_none()
             if updated_row is None:
                 raise ArtifactUpdateError()
 
@@ -674,11 +670,7 @@ async def reject_artifact(self, revision_id: uuid.UUID) -> ArtifactRevisionData:
             )
 
             result = await db_sess.execute(update_stmt)
-            updated_id = result.scalar_one_or_none()
-            if updated_id is None:
-                raise ArtifactUpdateError()
-
-            updated_row = await db_sess.get(ArtifactRevisionRow, updated_id)
+            updated_row = result.scalars().one_or_none()
             if updated_row is None:
                 raise ArtifactUpdateError()
 

tests/unit/manager/services/auth/test_authorize.py

@@ -57,13 +57,14 @@ def setup_successful_auth(mock_auth_repository, mock_hook_plugin_ctx):
     mock_hook_plugin_ctx.dispatch.return_value = HookResult(
         status=HookResults.PASSED, result=None, reason=None
     )
-    mock_auth_repository.check_credential_with_migration.return_value = {
-        "uuid": UUID("12345678-1234-5678-1234-567812345678"),
-        "email": "test@example.com",
-        "role": UserRole.USER,
-        "status": UserStatus.ACTIVE,
-        "password_changed_at": None,
-    }
+    mock_user = MagicMock()
+    mock_user.uuid = UUID("12345678-1234-5678-1234-567812345678")
+    mock_user.email = "test@example.com"
+    mock_user.role = UserRole.USER
+    mock_user.status = UserStatus.ACTIVE
+    mock_user.password_changed_at = None
+    mock_user.__getitem__ = lambda self, key: getattr(self, key)
+    mock_auth_repository.check_credential_with_migration.return_value = mock_user
     mock_user_row = MagicMock()
     mock_user_row.get_main_keypair_row.return_value = MagicMock(
         access_key="test_access_key",
@@ -149,14 +150,14 @@ async def test_authorize_with_hook_authorization(
         stoken=None,
     )
 
-    # Hook returns user data
-    hook_user = {
-        "uuid": UUID("87654321-4321-8765-4321-876543218765"),
-        "email": "hook@example.com",
-        "role": UserRole.ADMIN,
-        "status": UserStatus.ACTIVE,
-        "password_changed_at": None,
-    }
+    # Hook returns user data as MagicMock for attribute access
+    hook_user = MagicMock()
+    hook_user.uuid = UUID("87654321-4321-8765-4321-876543218765")
+    hook_user.email = "hook@example.com"
+    hook_user.role = UserRole.ADMIN
+    hook_user.status = UserStatus.ACTIVE
+    hook_user.password_changed_at = None
+    hook_user.__getitem__ = lambda self, key: getattr(self, key)
     mock_hook_plugin_ctx.dispatch.return_value = HookResult(
         status=HookResults.PASSED,
         result=hook_user,
@@ -177,7 +178,7 @@ async def test_authorize_with_hook_authorization(
     assert result.authorization_result is not None
     assert result.authorization_result.access_key == "hook_access_key"
     assert result.authorization_result.secret_key == "hook_secret_key"
-    assert result.authorization_result.user_id == hook_user["uuid"]
+    assert result.authorization_result.user_id == hook_user.uuid
     assert result.authorization_result.role == UserRole.ADMIN
 
 
@@ -202,13 +203,14 @@ async def test_authorize_with_password_expiry(
 
     # Setup expired password
     password_changed_at = datetime.now(tz=UTC) - timedelta(days=100)
-    mock_auth_repository.check_credential_with_migration.return_value = {
-        "uuid": UUID("12345678-1234-5678-1234-567812345678"),
-        "email": "expired@example.com",
-        "role": UserRole.USER,
-        "status": UserStatus.ACTIVE,
-        "password_changed_at": password_changed_at,
-    }
+    mock_user = MagicMock()
+    mock_user.uuid = UUID("12345678-1234-5678-1234-567812345678")
+    mock_user.email = "expired@example.com"
+    mock_user.role = UserRole.USER
+    mock_user.status = UserStatus.ACTIVE
+    mock_user.password_changed_at = password_changed_at
+    mock_user.__getitem__ = lambda self, key: getattr(self, key)
+    mock_auth_repository.check_credential_with_migration.return_value = mock_user
     mock_auth_repository.get_current_time.return_value = datetime.now(tz=UTC)
 
     mock_hook_plugin_ctx.dispatch.return_value = HookResult(
@@ -238,13 +240,14 @@ async def test_authorize_with_post_hook_response(
     )
 
     # Setup successful credential check
-    mock_auth_repository.check_credential_with_migration.return_value = {
-        "uuid": UUID("12345678-1234-5678-1234-567812345678"),
-        "email": "test@example.com",
-        "role": UserRole.USER,
-        "status": UserStatus.ACTIVE,
-        "password_changed_at": None,
-    }
+    mock_user = MagicMock()
+    mock_user.uuid = UUID("12345678-1234-5678-1234-567812345678")
+    mock_user.email = "test@example.com"
+    mock_user.role = UserRole.USER
+    mock_user.status = UserStatus.ACTIVE
+    mock_user.password_changed_at = None
+    mock_user.__getitem__ = lambda self, key: getattr(self, key)
+    mock_auth_repository.check_credential_with_migration.return_value = mock_user
 
     # Mock user row
     mock_user_row = MagicMock()

tests/unit/manager/services/scaling_group/test_scaling_group_service.py

@@ -434,12 +434,11 @@ async def test_check_scaling_group_raises_session_type_not_allowed(
     ) -> None:
         """Test that check_scaling_group raises ScalingGroupSessionTypeNotAllowed (422)
         when requesting BATCH session on INTERACTIVE-only scaling group"""
-        mock_sgroup = {
-            "name": "test-sgroup",
-            "scheduler_opts": ScalingGroupOpts(
-                allowed_session_types=[SessionTypes.INTERACTIVE],
-            ),
-        }
+        mock_sgroup = MagicMock()
+        mock_sgroup.name = "test-sgroup"
+        mock_sgroup.scheduler_opts = ScalingGroupOpts(
+            allowed_session_types=[SessionTypes.INTERACTIVE],
+        )
 
         with patch(
             "ai.backend.manager.registry.query_allowed_sgroups",
@@ -462,12 +461,11 @@ async def test_check_scaling_group_succeeds_with_allowed_session_type(
         mock_conn: MagicMock,
     ) -> None:
         """Test that check_scaling_group succeeds when session type is allowed"""
-        mock_sgroup = {
-            "name": "test-sgroup",
-            "scheduler_opts": ScalingGroupOpts(
-                allowed_session_types=[SessionTypes.INTERACTIVE],
-            ),
-        }
+        mock_sgroup = MagicMock()
+        mock_sgroup.name = "test-sgroup"
+        mock_sgroup.scheduler_opts = ScalingGroupOpts(
+            allowed_session_types=[SessionTypes.INTERACTIVE],
+        )
 
         with patch(
             "ai.backend.manager.registry.query_allowed_sgroups",

tests/unit/manager/services/session/test_session_service.py

@@ -389,7 +389,7 @@ async def test_empty_status_history(
         sample_session_id: SessionId,
         sample_access_key: AccessKey,
     ) -> None:
-        """Test getting empty status history"""
+        """Test getting empty status history returns empty dict when None"""
         mock_session = MagicMock()
         mock_session.id = sample_session_id
         mock_session.status_history = None
@@ -402,7 +402,7 @@ async def test_empty_status_history(
         result = await session_service.get_status_history(action)
 
         assert result.session_id == sample_session_id
-        assert result.status_history is None
+        assert result.status_history == {}
 
 
 # ==================== DestroySession Tests ====================
@@ -624,7 +624,7 @@ async def test_success(
         mock_kernel.image = "cr.backend.ai/stable/python:latest"
         mock_kernel.architecture = "x86_64"
         mock_kernel.registry = "cr.backend.ai"
-        mock_kernel.container_id = uuid4()
+        mock_kernel.container_id = str(uuid4())
         mock_kernel.occupied_slots = ResourceSlot({"cpu": 1, "mem": 1024})
         mock_kernel.occupied_shares = {}
 

tests/unit/manager/test_scaling_group_for_enqueue.py

@@ -14,28 +14,23 @@
 from ai.backend.manager.registry import check_scaling_group
 
 
+def _create_mock_sgroup(name: str, allowed_session_types: list[str]) -> MagicMock:
+    """Create a mock scaling group with proper attribute access."""
+    mock = MagicMock()
+    mock.name = name
+    mock.scheduler_opts = ScalingGroupOpts.from_json({
+        "allowed_session_types": allowed_session_types,
+    })
+    return mock
+
+
 @pytest.mark.asyncio
 @mock.patch("ai.backend.manager.registry.query_allowed_sgroups")
 async def test_allowed_session_types_check(mock_query):
     mock_query.return_value = [
-        {
-            "name": "a",
-            "scheduler_opts": ScalingGroupOpts().from_json({
-                "allowed_session_types": ["batch"],
-            }),
-        },
-        {
-            "name": "b",
-            "scheduler_opts": ScalingGroupOpts().from_json({
-                "allowed_session_types": ["interactive"],
-            }),
-        },
-        {
-            "name": "c",
-            "scheduler_opts": ScalingGroupOpts().from_json({
-                "allowed_session_types": ["batch", "interactive"],
-            }),
-        },
+        _create_mock_sgroup("a", ["batch"]),
+        _create_mock_sgroup("b", ["interactive"]),
+        _create_mock_sgroup("c", ["batch", "interactive"]),
     ]
     mock_conn = MagicMock()
     mock_sess_ctx = MagicMock()
@@ -126,12 +121,7 @@ async def test_allowed_session_types_check(mock_query):
     # No preferred scaling group with a non-empty list of allowed sgroups
 
     mock_query.return_value = [
-        {
-            "name": "a",
-            "scheduler_opts": ScalingGroupOpts.from_json({
-                "allowed_session_types": ["batch"],
-            }),
-        },
+        _create_mock_sgroup("a", ["batch"]),
     ]
 
     session_type = SessionTypes.BATCH