From 43c8d2509242d0d7b3b3c5f70275af3e07000856 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Fri, 26 Sep 2025 00:33:28 +0100
Subject: [PATCH] Fix tests

---
 .../unit/test_base/test_flow_parameter_specs.py  | 11 +++++++++--
 .../test_base/test_llm_service_parameters.py     | 16 ++++++++++++++--
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/tests/unit/test_base/test_flow_parameter_specs.py b/tests/unit/test_base/test_flow_parameter_specs.py
index 82d2c0db..d9c3e606 100644
--- a/tests/unit/test_base/test_flow_parameter_specs.py
+++ b/tests/unit/test_base/test_flow_parameter_specs.py
@@ -14,6 +14,7 @@ from trustgraph.base import ParameterSpec, ConsumerSpec, ProducerSpec
 def mock_async_processor_init(self, **params):
     """Mock AsyncProcessor.__init__ that properly initializes required attributes"""
     self.config_handlers = []
+    self.id = params.get('id', 'test-service')
 
 
 # Apply the mock globally for this test module
@@ -28,6 +29,7 @@ class TestFlowParameterSpecs(IsolatedAsyncioTestCase):
         # Arrange
         def mock_init(self, **kwargs):
             self.config_handlers = []
+            self.id = kwargs.get('id', 'test-service')
 
         config = {
             'id': 'test-flow-processor',
@@ -60,6 +62,7 @@ class TestFlowParameterSpecs(IsolatedAsyncioTestCase):
         # Arrange
         def mock_init(self, **kwargs):
             self.config_handlers = []
+            self.id = kwargs.get('id', 'test-service')
 
         config = {
             'id': 'test-flow-processor',
@@ -70,7 +73,7 @@ class TestFlowParameterSpecs(IsolatedAsyncioTestCase):
 
         # Create different spec types
         param_spec = ParameterSpec(name="model")
-        consumer_spec = ConsumerSpec(name="input")
+        consumer_spec = ConsumerSpec(name="input", schema=MagicMock(), handler=MagicMock())
         producer_spec = ProducerSpec(name="output")
 
         # Act
@@ -95,6 +98,7 @@ class TestFlowParameterSpecs(IsolatedAsyncioTestCase):
         # Arrange
         def mock_init(self, **kwargs):
             self.config_handlers = []
+            self.id = kwargs.get('id', 'test-service')
 
         config = {
             'id': 'test-flow-processor',
@@ -128,6 +132,7 @@ class TestFlowParameterSpecs(IsolatedAsyncioTestCase):
         # Arrange
         def mock_init(self, **kwargs):
             self.config_handlers = []
+            self.id = kwargs.get('id', 'test-service')
 
         config = {
             'id': 'test-flow-processor',
@@ -152,11 +157,12 @@ class TestFlowParameterSpecs(IsolatedAsyncioTestCase):
         assert len(param_specs) >= 1  # At least one should be registered
 
     @patch('trustgraph.base.flow_processor.Flow')
-    async def test_parameter_specs_available_to_flows(self, mock_async_init, mock_flow_class):
+    async def test_parameter_specs_available_to_flows(self, mock_flow_class):
         """Test that parameter specs are available when flows are created"""
         # Arrange
         def mock_init(self, **kwargs):
             self.config_handlers = []
+            self.id = kwargs.get('id', 'test-service')
 
         config = {
             'id': 'test-flow-processor',
@@ -197,6 +203,7 @@ class TestParameterSpecValidation(IsolatedAsyncioTestCase):
         # Arrange
         def mock_init(self, **kwargs):
             self.config_handlers = []
+            self.id = kwargs.get('id', 'test-service')
 
         config = {
             'id': 'test-flow-processor',
diff --git a/tests/unit/test_base/test_llm_service_parameters.py b/tests/unit/test_base/test_llm_service_parameters.py
index 6bfafe1d..895748fb 100644
--- a/tests/unit/test_base/test_llm_service_parameters.py
+++ b/tests/unit/test_base/test_llm_service_parameters.py
@@ -103,6 +103,10 @@ class TestLlmServiceParameters(IsolatedAsyncioTestCase):
 
         service = LlmService(**config)
 
+        # Mock the metrics
+        service.text_completion_model_metric = MagicMock()
+        service.text_completion_model_metric.labels.return_value.info = AsyncMock()
+
         # Mock the generate_content method to capture parameters
         service.generate_content = AsyncMock(return_value=LlmResult(
             text="test response",
@@ -130,7 +134,7 @@ class TestLlmServiceParameters(IsolatedAsyncioTestCase):
         }.get(param, f"mock-{param}")
 
         mock_producer = AsyncMock()
-        mock_flow.return_value = mock_producer  # flow("response") returns producer
+        mock_flow.producer = {"response": mock_producer}
 
         # Act
         await service.on_request(mock_message, mock_consumer, mock_flow)
@@ -166,6 +170,10 @@ class TestLlmServiceParameters(IsolatedAsyncioTestCase):
 
         service = LlmService(**config)
 
+        # Mock the metrics
+        service.text_completion_model_metric = MagicMock()
+        service.text_completion_model_metric.labels.return_value.info = AsyncMock()
+
         # Mock the generate_content method
         service.generate_content = AsyncMock(return_value=LlmResult(
             text="test response",
@@ -189,7 +197,7 @@ class TestLlmServiceParameters(IsolatedAsyncioTestCase):
         mock_flow.return_value = None  # Both parameters return None
 
         mock_producer = AsyncMock()
-        mock_flow.return_value = mock_producer
+        mock_flow.producer = {"response": mock_producer}
 
         # Act
         await service.on_request(mock_message, mock_consumer, mock_flow)
@@ -221,6 +229,10 @@ class TestLlmServiceParameters(IsolatedAsyncioTestCase):
 
         service = LlmService(**config)
 
+        # Mock the metrics
+        service.text_completion_model_metric = MagicMock()
+        service.text_completion_model_metric.labels.return_value.info = AsyncMock()
+
         # Mock generate_content to raise an exception
         service.generate_content = AsyncMock(side_effect=Exception("Test error"))