wip

Author: jopemachine

tests/component/session/test_session_lifecycle.py

@@ -23,6 +23,7 @@
     TransitSessionStatusRequest,
 )
 from ai.backend.common.dto.manager.session.response import (
+    DestroySessionResponse,
     GetSessionInfoResponse,
     GetStatusHistoryResponse,
     TransitSessionStatusResponse,
@@ -147,9 +148,11 @@ async def test_restart_running_session(
             session_seed.session_name,
             RestartSessionRequest(),
         )
-        # After restart, session should still be accessible
+        # After restart, session should still be accessible and running
         result = await admin_registry.session.get_info(session_seed.session_name)
         assert isinstance(result, GetSessionInfoResponse)
+        assert result.root["status"] == "RUNNING"
+        assert result.root["domainName"] == session_seed.domain_name
 
     async def test_restart_terminated_session_fails(
         self,
@@ -195,7 +198,8 @@ async def test_transit_single_session_status(
             TransitSessionStatusRequest(ids=[session_seed.session_id]),
         )
         assert isinstance(result, TransitSessionStatusResponse)
-        assert isinstance(result.session_status_map, dict)
+        assert session_seed.session_id in result.session_status_map
+        assert result.session_status_map[session_seed.session_id] == SessionStatus.RUNNING.name
 
     async def test_transit_multiple_sessions(
         self,
@@ -226,7 +230,8 @@ async def test_transit_multiple_sessions(
             ),
         )
         assert isinstance(result, TransitSessionStatusResponse)
-        assert isinstance(result.session_status_map, dict)
+        assert session_seed.session_id in result.session_status_map
+        assert terminated_session_seed.session_id in result.session_status_map
 
     async def test_user_cannot_transit_others_session(
         self,
@@ -265,6 +270,10 @@ async def test_rename_back_and_forth(
         result = await admin_registry.session.get_info(new_name)
         assert isinstance(result, GetSessionInfoResponse)
         assert result.root["status"] == "RUNNING"
+        assert result.root["domainName"] == session_seed.domain_name
+        # Old name should no longer resolve
+        with pytest.raises(NotFoundError):
+            await admin_registry.session.get_info(original_name)
 
         # Rename back to original name
         await admin_registry.session.rename(
@@ -274,6 +283,7 @@ async def test_rename_back_and_forth(
         result = await admin_registry.session.get_info(original_name)
         assert isinstance(result, GetSessionInfoResponse)
         assert result.root["status"] == "RUNNING"
+        assert result.root["domainName"] == session_seed.domain_name
 
     async def test_rename_to_same_name_is_rejected(
         self,
@@ -301,6 +311,10 @@ async def test_user_renames_own_session(
         result = await user_registry.session.get_info(new_name)
         assert isinstance(result, GetSessionInfoResponse)
         assert result.root["status"] == "RUNNING"
+        assert result.root["domainName"] == user_session_seed.domain_name
+        # Old name should no longer resolve
+        with pytest.raises(NotFoundError):
+            await user_registry.session.get_info(user_session_seed.session_name)
 
 
 class TestSessionPermissions:
@@ -315,6 +329,7 @@ async def test_admin_gets_own_session_info(
         result = await admin_registry.session.get_info(session_seed.session_name)
         assert isinstance(result, GetSessionInfoResponse)
         assert result.root["status"] == "RUNNING"
+        assert result.root["domainName"] == session_seed.domain_name
 
     async def test_user_gets_own_session_info(
         self,
@@ -325,6 +340,8 @@ async def test_user_gets_own_session_info(
         result = await user_registry.session.get_info(user_session_seed.session_name)
         assert isinstance(result, GetSessionInfoResponse)
         assert result.root["status"] == "RUNNING"
+        assert result.root["domainName"] == user_session_seed.domain_name
+        assert result.root["userId"] == str(user_session_seed.user_uuid)
 
     async def test_user_cannot_access_admin_session(
         self,
@@ -419,6 +436,7 @@ async def test_user_gets_own_session_status_history(
         assert isinstance(result, GetStatusHistoryResponse)
         history = result.root
         assert isinstance(history, dict)
+        assert "PENDING" in history
         assert "RUNNING" in history
 
     async def test_user_cannot_access_admin_session_status_history(
@@ -446,7 +464,8 @@ async def test_destroy_already_terminated_session_succeeds(
             terminated_session_seed.session_name,
             DestroySessionRequest(forced=True),
         )
-        assert result is not None
+        assert isinstance(result, DestroySessionResponse)
+        assert isinstance(result.root, dict)
 
     async def test_user_destroys_own_session(
         self,
@@ -460,5 +479,8 @@ async def test_user_destroys_own_session(
             user_session_seed.session_name,
             DestroySessionRequest(forced=True),
         )
-        # Verify the session status changed
-        assert result is not None
+        assert isinstance(result, DestroySessionResponse)
+        assert isinstance(result.root, dict)
+        # Session should no longer be accessible after destroy
+        with pytest.raises((NotFoundError, BackendAPIError)):
+            await user_registry.session.get_info(user_session_seed.session_name)