From f763f32b04d6b91b338f689bac07f7c91fcc8025 Mon Sep 17 00:00:00 2001
From: adavyas <adavyasharma@gmail.com>
Date: Sat, 14 Mar 2026 13:12:58 -0700
Subject: [PATCH] Fix custom instruction test typing

---
 tests/deriver/test_deriver_processing.py |  5 +++--
 tests/utils/test_config_helpers.py       | 13 ++++++++-----
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/tests/deriver/test_deriver_processing.py b/tests/deriver/test_deriver_processing.py
index fb44989d..a2456e61 100644
--- a/tests/deriver/test_deriver_processing.py
+++ b/tests/deriver/test_deriver_processing.py
@@ -1,6 +1,5 @@
 import signal
 from datetime import datetime, timezone
-from types import SimpleNamespace
 from typing import Any
 from unittest.mock import AsyncMock
 
@@ -234,12 +233,14 @@ class TestCustomInstructions:
             AsyncMock(),
         )
 
-        message = SimpleNamespace(
+        message = models.Message(
             id=1,
+            public_id="msg_1",
             content="I like tea.",
             created_at=datetime.now(timezone.utc),
             peer_name="alice",
             token_count=4,
+            seq_in_session=1,
             session_name="session-1",
             workspace_name="workspace-1",
         )
diff --git a/tests/utils/test_config_helpers.py b/tests/utils/test_config_helpers.py
index 26d6ed13..e4c4b56f 100644
--- a/tests/utils/test_config_helpers.py
+++ b/tests/utils/test_config_helpers.py
@@ -1,12 +1,12 @@
-from types import SimpleNamespace
-
+from src import models
 from src.schemas import MessageConfiguration, ReasoningConfiguration
 from src.utils.config_helpers import get_configuration
 
 
 class TestGetConfiguration:
     def test_preserves_workspace_custom_instructions(self) -> None:
-        workspace = SimpleNamespace(
+        workspace = models.Workspace(
+            name="workspace-1",
             configuration={
                 "reasoning": {
                     "enabled": True,
@@ -21,10 +21,13 @@ class TestGetConfiguration:
         assert config.reasoning.custom_instructions == "Focus on durable preferences."
 
     def test_message_custom_instructions_override_session_and_workspace(self) -> None:
-        workspace = SimpleNamespace(
+        workspace = models.Workspace(
+            name="workspace-1",
             configuration={"reasoning": {"custom_instructions": "workspace scope"}}
         )
-        session = SimpleNamespace(
+        session = models.Session(
+            name="session-1",
+            workspace_name="workspace-1",
             configuration={"reasoning": {"custom_instructions": "session scope"}}
         )
         message = MessageConfiguration(