feat: Add session history retriever test

Author: seedspirit

changes/4697.feature.md

@@ -0,0 +1 @@
+Add session status history retrieval test validating history not empty and contains valid statuses

src/ai/backend/test/contexts/session.py

@@ -1,7 +1,7 @@
 from typing import override
-from uuid import UUID
 
 from ai.backend.test.contexts.context import BaseTestContext, ContextName
+from ai.backend.test.data.session import CreatedSessionMeta
 from ai.backend.test.tester.dependency import (
     BatchSessionDep,
     ClusterDep,
@@ -45,7 +45,7 @@ def name(cls) -> ContextName:
         return ContextName.CREATED_SESSION_TEMPLATE_ID
 
 
-class CreatedSessionIDContext(BaseTestContext[UUID]):
+class CreatedSessionMetaContext(BaseTestContext[CreatedSessionMeta]):
     @override
     @classmethod
     def name(cls) -> ContextName:

src/ai/backend/test/data/session.py

@@ -0,0 +1,8 @@
+import uuid
+from dataclasses import dataclass
+
+
+@dataclass
+class CreatedSessionMeta:
+    id: uuid.UUID
+    name: str

src/ai/backend/test/templates/session/batch_session.py

@@ -9,13 +9,14 @@
 from ai.backend.test.contexts.session import (
     BatchSessionContext,
     ClusterContext,
-    CreatedSessionIDContext,
+    CreatedSessionMetaContext,
     SessionContext,
 )
 from ai.backend.test.contexts.sse import (
     SSEContext,
 )
 from ai.backend.test.contexts.tester import TestSpecMetaContext
+from ai.backend.test.data.session import CreatedSessionMeta
 from ai.backend.test.templates.session.utils import verify_session_events
 from ai.backend.test.templates.template import (
     WrapperTestTemplate,
@@ -78,5 +79,7 @@ async def _context(self) -> AsyncIterator[None]:
         session_name = f"test_session_{str(test_id)}"
 
         session_id = await self._verify_session_creation(client_session, session_name)
-        with CreatedSessionIDContext.with_current(session_id):
+        with CreatedSessionMetaContext.with_current(
+            CreatedSessionMeta(id=session_id, name=session_name)
+        ):
             yield

src/ai/backend/test/templates/session/interactive_session.py

@@ -8,13 +8,14 @@
 from ai.backend.test.contexts.image import ImageContext
 from ai.backend.test.contexts.session import (
     ClusterContext,
-    CreatedSessionIDContext,
+    CreatedSessionMetaContext,
     SessionContext,
 )
 from ai.backend.test.contexts.sse import (
     SSEContext,
 )
 from ai.backend.test.contexts.tester import TestSpecMetaContext
+from ai.backend.test.data.session import CreatedSessionMeta
 from ai.backend.test.templates.session.utils import (
     verify_session_events,
 )
@@ -100,7 +101,9 @@ async def _context(self) -> AsyncIterator[None]:
         session_id = None
         try:
             session_id = await self._verify_session_creation(client_session, session_name)
-            with CreatedSessionIDContext.with_current(session_id):
+            with CreatedSessionMetaContext.with_current(
+                CreatedSessionMeta(id=session_id, name=session_name)
+            ):
                 yield
         finally:
             if session_id:

src/ai/backend/test/templates/session/session_template.py

@@ -10,14 +10,15 @@
 from ai.backend.test.contexts.session import (
     BatchSessionContext,
     ClusterContext,
-    CreatedSessionIDContext,
+    CreatedSessionMetaContext,
     CreatedSessionTemplateIDContext,
     SessionContext,
 )
 from ai.backend.test.contexts.sse import (
     SSEContext,
 )
 from ai.backend.test.contexts.tester import TestSpecMetaContext
+from ai.backend.test.data.session import CreatedSessionMeta
 from ai.backend.test.templates.session.utils import (
     verify_session_events,
 )
@@ -151,7 +152,9 @@ async def _context(self) -> AsyncIterator[None]:
         session_name = f"test_session_{str(test_id)}"
 
         session_id = await self._verify_session_creation(client_session, session_name)
-        with CreatedSessionIDContext.with_current(session_id):
+        with CreatedSessionMetaContext.with_current(
+            CreatedSessionMeta(id=session_id, name=session_name)
+        ):
             yield
 
 
@@ -228,7 +231,9 @@ async def _context(self) -> AsyncIterator[None]:
         session_id = None
         try:
             session_id = await self._verify_session_creation(client_session, session_name)
-            with CreatedSessionIDContext.with_current(session_id):
+            with CreatedSessionMetaContext.with_current(
+                CreatedSessionMeta(id=session_id, name=session_name)
+            ):
                 yield
         finally:
             if session_id:

src/ai/backend/test/testcases/session/execution.py

@@ -1,12 +1,13 @@
 from ai.backend.test.contexts.client_session import ClientSessionContext
-from ai.backend.test.contexts.session import CodeExecutionContext, CreatedSessionIDContext
+from ai.backend.test.contexts.session import CodeExecutionContext, CreatedSessionMetaContext
 from ai.backend.test.templates.template import TestCode
 
 
 class InteractiveSessionExecuteCodeSuccess(TestCode):
     async def test(self) -> None:
         client_session = ClientSessionContext.current()
-        session_id = CreatedSessionIDContext.current()
+        session_meta = CreatedSessionMetaContext.current()
+        session_id = session_meta.id
         code_dep = CodeExecutionContext.current()
 
         result = await client_session.ComputeSession(str(session_id)).execute(
@@ -41,7 +42,8 @@ async def test(self) -> None:
 class InteractiveSessionExecuteCodeFailureWrongCommand(TestCode):
     async def test(self) -> None:
         client_session = ClientSessionContext.current()
-        session_id = CreatedSessionIDContext.current()
+        session_meta = CreatedSessionMetaContext.current()
+        session_id = session_meta.id
         WRONG_CMD = "some wrong command !@#"
 
         result = await client_session.ComputeSession(str(session_id)).execute(

src/ai/backend/test/testcases/session/status_history_retriever.py

@@ -0,0 +1,38 @@
+from typing import override
+
+from ai.backend.test.contexts.client_session import ClientSessionContext
+from ai.backend.test.contexts.session import CreatedSessionMetaContext
+from ai.backend.test.templates.template import TestCode
+
+
+class StatusHistoryRetriever(TestCode):
+    @override
+    async def test(self) -> None:
+        session = ClientSessionContext.current()
+        session_meta = CreatedSessionMetaContext.current()
+        session_name = session_meta.name
+
+        expected_status = set([
+            "PENDING",
+            "SCHEDULED",
+            "PREPARING",
+            "PULLING",
+            "PREPARED",
+            "CREATING",
+            "RUNNING",
+            "RESTARTING",
+            "RUNNING_DEGRADED",
+            "TERMINATING",
+            "TERMINATED",
+            "ERROR",
+            "CANCELED",
+        ])
+
+        result = await session.ComputeSession(name=session_name).get_status_history()
+        assert result["result"] != "", "Status history should not be empty"
+        # NOTE: SessionStatus enum is not available due to BUILD rules
+        # NOTE: src/ai/backend/test/BUILD[//src/ai/backend/**] -> src/ai/backend/manager/BUILD[!*] : DENY
+        for status_key, status_value in result["result"].items():
+            assert status_key in expected_status, (
+                f"Status {status_key} is not a valid session status"
+            )

src/ai/backend/test/testcases/session/testspecs.py

@@ -10,6 +10,7 @@
     InteractiveSessionFromTemplateTemplate,
     SessionTemplateTemplate,
 )
+from ai.backend.test.templates.template import BasicTestTemplate, NopTestCode
 from ai.backend.test.testcases.session.creation_failure_low_resources import (
     SessionCreationFailureLowResources,
 )
@@ -23,11 +24,10 @@
     InteractiveSessionExecuteCodeFailureWrongCommand,
     InteractiveSessionExecuteCodeSuccess,
 )
+from ai.backend.test.testcases.session.status_history_retriever import StatusHistoryRetriever
 from ai.backend.test.testcases.spec_manager import TestSpec, TestTag
 from ai.backend.test.tester.dependency import ClusterDep, CodeExecutionDep
 
-from ...templates.template import BasicTestTemplate, NopTestCode
-
 BATCH_SESSION_TEST_SPECS = {
     "creation_batch_session_success": TestSpec(
         name="creation_batch_session_success",
@@ -265,9 +265,31 @@
     ),
 }
 
+SESSION_INFO_RETRIEVER_TEST_SPECS = {
+    "session_status_history": TestSpec(
+        name="session_status_history",
+        description=textwrap.dedent("""
+            Tests retrieval of session status history
+            Validate that the status history is not empty and contains valid statuses
+        """),
+        tags={TestTag.SESSION, TestTag.MANAGER},
+        template=BasicTestTemplate(
+            testcode=StatusHistoryRetriever(),
+        ).with_wrappers(KeypairAuthTemplate, InteractiveSessionTemplate),
+        parametrizes={
+            ContextName.CLUSTER_CONFIG: [
+                ClusterDep(
+                    cluster_mode=ClusterMode.SINGLE_NODE,
+                    cluster_size=1,
+                ),
+            ]
+        },
+    ),
+}
 
 SESSION_TEST_SPECS = {
-    **BATCH_SESSION_TEST_SPECS,
-    **INTERACTIVE_SESSION_TEST_SPECS,
-    **SESSION_TEMPLATE_TEST_SPECS,
+    # **BATCH_SESSION_TEST_SPECS,
+    # **INTERACTIVE_SESSION_TEST_SPECS,
+    # **SESSION_TEMPLATE_TEST_SPECS,
+    **SESSION_INFO_RETRIEVER_TEST_SPECS
 }