Unit & Integration tests for State Manager

.github/workflows/test-state-manager.yml

@@ -0,0 +1,67 @@
+name: State Manager Unit Tests
+
+on:
+  push:
+    branches: [main]
+    paths:
+      - 'state-manager/**'
+  pull_request:
+    branches: [main]
+    paths:
+      - 'state-manager/**'
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    services:
+      mongodb:
+        image: mongo:7
+        ports:
+          - 27017:27017
+        options: >-
+          --health-cmd "mongosh --eval 'db.runCommand(\"ping\")'"
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v2
+        with:
+          cache: true
+
+      - name: Install dev dependencies with uv
+        working-directory: state-manager
+        run: |
+          uv sync --group dev
+
+      - name: Run unit tests with pytest and coverage
+        working-directory: state-manager
+        run: |
+          uv run pytest tests/unit/ --cov=app --cov-report=xml --cov-report=term-missing -v --junitxml=pytest-report.xml
+
+      - name: Upload coverage reports to Codecov
+        uses: codecov/codecov-action@v5
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          slug: exospherehost/exospherehost
+          files: state-manager/coverage.xml
+          flags: state-manager-unittests
+          name: state-manager-coverage-report
+          fail_ci_if_error: true
+
+      - name: Upload test results
+        uses: actions/upload-artifact@v4
+        if: always()
+        with:
+          name: state-manager-test-results
+          path: state-manager/pytest-report.xml
+          retention-days: 30

.vscode/settings.json

@@ -12,5 +12,10 @@
       "tag:yaml.org,2002:python/object/apply:pymdownx.slugs.slugify mapping"
     ],
     "python.analysis.typeCheckingMode": "basic",
-    "python.analysis.autoImportCompletions": true
+    "python.analysis.autoImportCompletions": true,
+    "python.testing.pytestArgs": [
+      "state-manager"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
 }

state-manager/app/controller/create_states.py

@@ -1,14 +1,15 @@
 from fastapi import HTTPException
 
 from app.singletons.logs_manager import LogsManager
-from app.models.create_models import CreateRequestModel, CreateResponseModel, ResponseStateModel
+from app.models.create_models import CreateRequestModel, CreateResponseModel, ResponseStateModel, TriggerGraphRequestModel, TriggerGraphResponseModel
 from app.models.state_status_enum import StateStatusEnum
 from app.models.db.state import State
 from app.models.db.graph_template_model import GraphTemplate
 from app.models.node_template_model import NodeTemplate
 
 from beanie.operators import In
-from bson import ObjectId
+from beanie import PydanticObjectId
+import uuid
 
 logger = LogsManager().get_logger()
 
@@ -20,6 +21,33 @@ def get_node_template(graph_template: GraphTemplate, identifier: str) -> NodeTem
     return node
 
 
+async def trigger_graph(namespace_name: str, graph_name: str, body: TriggerGraphRequestModel, x_exosphere_request_id: str) -> TriggerGraphResponseModel:
+    try:
+        # Generate a new run ID for this graph execution
+        run_id = str(uuid.uuid4())
+        logger.info(f"Triggering graph {graph_name} with run_id {run_id}", x_exosphere_request_id=x_exosphere_request_id)
+
+        # Create a CreateRequestModel with the generated run_id
+        create_request = CreateRequestModel(
+            run_id=run_id,
+            states=body.states
+        )
+
+        # Call the existing create_states function
+        create_response = await create_states(namespace_name, graph_name, create_request, x_exosphere_request_id)
+
+        # Return the trigger response with the generated run_id
+        return TriggerGraphResponseModel(
+            run_id=run_id,
+            status=create_response.status,
+            states=create_response.states
+        )
+
+    except Exception as e:
+        logger.error(f"Error triggering graph {graph_name} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+        raise e
+
+
 async def create_states(namespace_name: str, graph_name: str, body: CreateRequestModel, x_exosphere_request_id: str) -> CreateResponseModel:
     try:
         states = []
@@ -39,6 +67,7 @@ async def create_states(namespace_name: str, graph_name: str, body: CreateReques
                     node_name=node_template.node_name,
                     namespace_name=node_template.namespace,
                     graph_name=graph_name,
+                    run_id=body.run_id,
                     status=StateStatusEnum.CREATED,
                     inputs=state.inputs,
                     outputs={},
@@ -51,12 +80,12 @@ async def create_states(namespace_name: str, graph_name: str, body: CreateReques
         logger.info(f"Created states: {inserted_states.inserted_ids}", x_exosphere_request_id=x_exosphere_request_id)
 
         newStates = await State.find(
-            In(State.id, [ObjectId(id) for id in inserted_states.inserted_ids])
+            In(State.id, [PydanticObjectId(id) for id in inserted_states.inserted_ids])
         ).to_list()
 
         return CreateResponseModel(
             status=StateStatusEnum.CREATED,
-            states=[ResponseStateModel(state_id=str(state.id), identifier=state.identifier, node_name=state.node_name, graph_name=state.graph_name, inputs=state.inputs, created_at=state.created_at) for state in newStates]
+            states=[ResponseStateModel(state_id=str(state.id), identifier=state.identifier, node_name=state.node_name, graph_name=state.graph_name, run_id=state.run_id, inputs=state.inputs, created_at=state.created_at) for state in newStates]
         )
 
     except Exception as e:

state-manager/app/controller/errored_state.py

@@ -1,14 +1,14 @@
 from app.models.errored_models import ErroredRequestModel, ErroredResponseModel
-from bson import ObjectId
 from fastapi import HTTPException, status
+from beanie import PydanticObjectId
 
 from app.models.db.state import State
 from app.models.state_status_enum import StateStatusEnum
 from app.singletons.logs_manager import LogsManager
 
 logger = LogsManager().get_logger()
 
-async def errored_state(namespace_name: str, state_id: ObjectId, body: ErroredRequestModel, x_exosphere_request_id: str) -> ErroredResponseModel:
+async def errored_state(namespace_name: str, state_id: PydanticObjectId, body: ErroredRequestModel, x_exosphere_request_id: str) -> ErroredResponseModel:
 
     try:
         logger.info(f"Errored state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
@@ -23,9 +23,9 @@ async def errored_state(namespace_name: str, state_id: ObjectId, body: ErroredRe
         if state.status == StateStatusEnum.EXECUTED:
             raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="State is already executed")
 
-        await State.find_one(State.id == state_id).set(
-            {"status": StateStatusEnum.ERRORED, "error": body.error}
-        )
+        state.status = StateStatusEnum.ERRORED
+        state.error = body.error
+        await state.save()
 
         return ErroredResponseModel(status=StateStatusEnum.ERRORED)
 

state-manager/app/controller/executed_state.py

@@ -1,5 +1,6 @@
+from beanie import PydanticObjectId
 from app.models.executed_models import ExecutedRequestModel, ExecutedResponseModel
-from bson import ObjectId
+
 from fastapi import HTTPException, status, BackgroundTasks
 
 from app.models.db.state import State
@@ -9,7 +10,7 @@
 
 logger = LogsManager().get_logger()
 
-async def executed_state(namespace_name: str, state_id: ObjectId, body: ExecutedRequestModel, x_exosphere_request_id: str, background_tasks: BackgroundTasks) -> ExecutedResponseModel:
+async def executed_state(namespace_name: str, state_id: PydanticObjectId, body: ExecutedRequestModel, x_exosphere_request_id: str, background_tasks: BackgroundTasks) -> ExecutedResponseModel:
 
     try:
         logger.info(f"Executed state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
@@ -22,17 +23,18 @@ async def executed_state(namespace_name: str, state_id: ObjectId, body: Executed
             raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="State is not queued")
 
         if len(body.outputs) == 0:
-            await State.find_one(State.id == state_id).set(
-                {"status": StateStatusEnum.EXECUTED, "outputs": {}, "parents": {**state.parents, state.identifier: state.id}}
-            )
+            state.status = StateStatusEnum.EXECUTED
+            state.outputs = {}
+            state.parents = {**state.parents, state.identifier: state.id}
+            await state.save()
 
             background_tasks.add_task(create_next_state, state)
 
-        else:
-
+        else:            
             state.outputs = body.outputs[0]
             state.status = StateStatusEnum.EXECUTED
             state.parents = {**state.parents, state.identifier: state.id}
+
             await state.save()
 
             background_tasks.add_task(create_next_state, state)
@@ -43,6 +45,7 @@ async def executed_state(namespace_name: str, state_id: ObjectId, body: Executed
                     namespace_name=state.namespace_name,
                     identifier=state.identifier,
                     graph_name=state.graph_name,
+                    run_id=state.run_id,
                     status=StateStatusEnum.CREATED,
                     inputs=state.inputs,
                     outputs=output,

state-manager/app/controller/get_current_states.py

@@ -0,0 +1,37 @@
+"""
+Controller for fetching current states in a namespace
+"""
+from typing import List
+
+from ..models.db.state import State
+from ..singletons.logs_manager import LogsManager
+
+
+async def get_current_states(namespace: str, request_id: str) -> List[State]:
+    """
+    Get all current states in a namespace
+
+    Args:
+        namespace: The namespace to search in
+        request_id: Request ID for logging
+
+    Returns:
+        List of all states in the namespace
+    """
+    logger = LogsManager().get_logger()
+
+    try:
+        logger.info(f"Fetching current states for namespace: {namespace}", x_exosphere_request_id=request_id)
+
+        # Find all states in the namespace
+        states = await State.find(
+            State.namespace_name == namespace
+        ).to_list()
+
+        logger.info(f"Found {len(states)} states for namespace: {namespace}", x_exosphere_request_id=request_id)
+
+        return states
+
+    except Exception as e:
+        logger.error(f"Error fetching current states for namespace {namespace}: {str(e)}", x_exosphere_request_id=request_id)
+        raise

state-manager/app/controller/get_states_by_run_id.py

@@ -0,0 +1,39 @@
+"""
+Controller for fetching states by run ID
+"""
+from typing import List
+
+from ..models.db.state import State
+from ..singletons.logs_manager import LogsManager
+
+
+async def get_states_by_run_id(namespace: str, run_id: str, request_id: str) -> List[State]:
+    """
+    Get all states for a given run ID in a namespace
+
+    Args:
+        namespace: The namespace to search in
+        run_id: The run ID to filter by
+        request_id: Request ID for logging
+
+    Returns:
+        List of states for the given run ID
+    """
+    logger = LogsManager().get_logger()
+
+    try:
+        logger.info(f"Fetching states for run ID: {run_id} in namespace: {namespace}", x_exosphere_request_id=request_id)
+
+        # Find all states for the run ID in the namespace
+        states = await State.find(
+            State.run_id == run_id,
+            State.namespace_name == namespace
+        ).to_list()
+
+        logger.info(f"Found {len(states)} states for run ID: {run_id}", x_exosphere_request_id=request_id)
+
+        return states
+
+    except Exception as e:
+        logger.error(f"Error fetching states for run ID {run_id} in namespace {namespace}: {str(e)}", x_exosphere_request_id=request_id)
+        raise

state-manager/app/controller/list_graph_templates.py

@@ -0,0 +1,37 @@
+"""
+Controller for listing graph templates by namespace
+"""
+from typing import List
+
+from ..models.db.graph_template_model import GraphTemplate
+from ..singletons.logs_manager import LogsManager
+
+
+async def list_graph_templates(namespace: str, request_id: str) -> List[GraphTemplate]:
+    """
+    List all graph templates for a given namespace
+
+    Args:
+        namespace: The namespace to list graph templates for
+        request_id: Request ID for logging
+
+    Returns:
+        List of graph templates
+    """
+    logger = LogsManager().get_logger()
+
+    try:
+        logger.info(f"Listing graph templates for namespace: {namespace}", x_exosphere_request_id=request_id)
+
+        # Find all graph templates for the namespace
+        templates = await GraphTemplate.find(
+            GraphTemplate.namespace == namespace
+        ).to_list()
+
+        logger.info(f"Found {len(templates)} graph templates for namespace: {namespace}", x_exosphere_request_id=request_id)
+
+        return templates
+
+    except Exception as e:
+        logger.error(f"Error listing graph templates for namespace {namespace}: {str(e)}", x_exosphere_request_id=request_id)
+        raise

state-manager/app/controller/list_registered_nodes.py

@@ -0,0 +1,37 @@
+"""
+Controller for listing registered nodes by namespace
+"""
+from typing import List
+
+from ..models.db.registered_node import RegisteredNode
+from ..singletons.logs_manager import LogsManager
+
+
+async def list_registered_nodes(namespace: str, request_id: str) -> List[RegisteredNode]:
+    """
+    List all registered nodes for a given namespace
+
+    Args:
+        namespace: The namespace to list nodes for
+        request_id: Request ID for logging
+
+    Returns:
+        List of registered nodes
+    """
+    logger = LogsManager().get_logger()
+
+    try:
+        logger.info(f"Listing registered nodes for namespace: {namespace}", x_exosphere_request_id=request_id)
+
+        # Find all registered nodes for the namespace
+        nodes = await RegisteredNode.find(
+            RegisteredNode.namespace == namespace
+        ).to_list()
+
+        logger.info(f"Found {len(nodes)} registered nodes for namespace: {namespace}", x_exosphere_request_id=request_id)
+
+        return nodes
+
+    except Exception as e:
+        logger.error(f"Error listing registered nodes for namespace {namespace}: {str(e)}", x_exosphere_request_id=request_id)
+        raise