From 656ca430b95ac56f3810bfebf776a791196ddf52 Mon Sep 17 00:00:00 2001
From: cybermaggedon <cybermaggedon@gmail.com>
Date: Wed, 1 Jul 2026 15:40:23 +0100
Subject: [PATCH] fix: wire variant into text-completion integration test mocks
 (#1008)

Tests using MagicMock processors need the variant, thinking mode,
and _build_kwargs/_extract_content methods bound to work with the
new variant-based API kwargs construction.
---
 .../test_text_completion_integration.py           | 15 ++++++++++++++-
 .../test_text_completion_streaming_integration.py | 12 ++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/tests/integration/test_text_completion_integration.py b/tests/integration/test_text_completion_integration.py
index 6615bf84..521f7d74 100644
--- a/tests/integration/test_text_completion_integration.py
+++ b/tests/integration/test_text_completion_integration.py
@@ -15,11 +15,20 @@ from openai.types.chat.chat_completion import Choice
 from openai.types.completion_usage import CompletionUsage
 
 from trustgraph.model.text_completion.openai.llm import Processor
+from trustgraph.model.text_completion.openai.variants import get_variant
 from trustgraph.exceptions import TooManyRequests
 from trustgraph.base import LlmResult
 from trustgraph.schema import TextCompletionRequest, TextCompletionResponse, Error
 
 
+def _wire_variant(processor):
+    """Attach variant methods to a MagicMock processor."""
+    processor.variant = get_variant("openai")
+    processor.thinking = "off"
+    processor._build_kwargs = Processor._build_kwargs.__get__(processor, Processor)
+    processor._extract_content = Processor._extract_content.__get__(processor, Processor)
+
+
 @pytest.mark.integration
 class TestTextCompletionIntegration:
     """Integration tests for OpenAI text completion service coordination"""
@@ -66,6 +75,7 @@ class TestTextCompletionIntegration:
 
         # Add the actual generate_content method from Processor class
         processor.generate_content = Processor.generate_content.__get__(processor, Processor)
+        _wire_variant(processor)
 
         return processor
 
@@ -119,6 +129,7 @@ class TestTextCompletionIntegration:
 
             # Add the actual generate_content method
             processor.generate_content = Processor.generate_content.__get__(processor, Processor)
+            _wire_variant(processor)
 
             # Act
             result = await processor.generate_content("System prompt", "User prompt")
@@ -129,7 +140,7 @@ class TestTextCompletionIntegration:
             assert result.in_token == 50
             assert result.out_token == 100
             # Note: result.model comes from mock response, not processor config
-            
+
             # Verify configuration was applied
             call_args = mock_openai_client.chat.completions.create.call_args
             assert call_args.kwargs['model'] == config['model']
@@ -247,6 +258,7 @@ class TestTextCompletionIntegration:
             processor.max_output = processor_config["max_output"]
             processor.openai = mock_openai_client
             processor.generate_content = Processor.generate_content.__get__(processor, Processor)
+            _wire_variant(processor)
             processors.append(processor)
 
         # Simulate multiple concurrent requests
@@ -354,6 +366,7 @@ class TestTextCompletionIntegration:
         processor.max_output = 2048
         processor.openai = mock_openai_client
         processor.generate_content = Processor.generate_content.__get__(processor, Processor)
+        _wire_variant(processor)
 
         # Act
         await processor.generate_content("System prompt", "User prompt")
diff --git a/tests/integration/test_text_completion_streaming_integration.py b/tests/integration/test_text_completion_streaming_integration.py
index 6968affa..caa3ec9c 100644
--- a/tests/integration/test_text_completion_streaming_integration.py
+++ b/tests/integration/test_text_completion_streaming_integration.py
@@ -11,6 +11,7 @@ from openai.types.chat import ChatCompletionChunk
 from openai.types.chat.chat_completion_chunk import Choice as StreamChoice, ChoiceDelta
 
 from trustgraph.model.text_completion.openai.llm import Processor
+from trustgraph.model.text_completion.openai.variants import get_variant
 from trustgraph.base import LlmChunk
 from tests.utils.streaming_assertions import (
     assert_streaming_chunks_valid,
@@ -18,6 +19,14 @@ from tests.utils.streaming_assertions import (
 )
 
 
+def _wire_variant(processor):
+    """Attach variant methods to a MagicMock processor."""
+    processor.variant = get_variant("openai")
+    processor.thinking = "off"
+    processor._build_kwargs = Processor._build_kwargs.__get__(processor, Processor)
+    processor._extract_content = Processor._extract_content.__get__(processor, Processor)
+
+
 @pytest.mark.integration
 class TestTextCompletionStreaming:
     """Integration tests for Text Completion streaming"""
@@ -69,6 +78,7 @@ class TestTextCompletionStreaming:
         processor.generate_content_stream = Processor.generate_content_stream.__get__(
             processor, Processor
         )
+        _wire_variant(processor)
 
         return processor
 
@@ -190,6 +200,7 @@ class TestTextCompletionStreaming:
         processor.generate_content_stream = Processor.generate_content_stream.__get__(
             processor, Processor
         )
+        _wire_variant(processor)
 
         # Act
         chunks = []
@@ -223,6 +234,7 @@ class TestTextCompletionStreaming:
             processor.generate_content_stream = Processor.generate_content_stream.__get__(
                 processor, Processor
             )
+            _wire_variant(processor)
 
             # Act
             chunks = []