✅ Update tests for breaking API changes

Signed-off-by: Evaline Ju <[email protected]>

Author: evaline-ju

tests/generative_detectors/test_base.py

@@ -1,7 +1,7 @@
 # Standard
-from dataclasses import dataclass
-from typing import Optional
-from unittest.mock import patch
+from dataclasses import dataclass, field
+from typing import Any, Optional
+from unittest.mock import MagicMock, patch
 import asyncio
 
 # Third Party
@@ -49,17 +49,23 @@ class MockHFConfig:
 @dataclass
 class MockModelConfig:
     task = "generate"
+    runner_type = "generate"
     tokenizer = MODEL_NAME
     trust_remote_code = False
     tokenizer_mode = "auto"
     max_model_len = 100
     tokenizer_revision = None
-    embedding_mode = False
     multimodal_config = MultiModalConfig()
-    diff_sampling_param: Optional[dict] = None
     hf_config = MockHFConfig()
     logits_processor_pattern = None
+    logits_processors: list[str] | None = None
+    diff_sampling_param: dict | None = None
     allowed_local_media_path: str = ""
+    allowed_media_domains: list[str] | None = None
+    encoder_config = None
+    generation_config: str = "auto"
+    media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
+    skip_tokenizer_init = False
 
     def get_diff_sampling_param(self):
         return self.diff_sampling_param or {}
@@ -78,18 +84,18 @@ async def _async_serving_detection_completion_init():
     """Initialize a chat completion base with string templates"""
     engine = MockEngine()
     engine.errored = False
-    model_config = await engine.get_model_config()
+    engine.model_config = MockModelConfig()
+    engine.input_processor = MagicMock()
+    engine.io_processor = MagicMock()
     models = OpenAIServingModels(
         engine_client=engine,
-        model_config=model_config,
         base_model_paths=BASE_MODEL_PATHS,
     )
 
     detection_completion = ChatCompletionDetectionBase(
         task_template="hello {{user_text}}",
         output_template="bye {{text}}",
         engine_client=engine,
-        model_config=model_config,
         models=models,
         response_role="assistant",
         chat_template=CHAT_TEMPLATE,

tests/generative_detectors/test_granite_guardian.py

@@ -1,8 +1,8 @@
 # Standard
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from http import HTTPStatus
-from typing import Optional
-from unittest.mock import patch
+from typing import Any, Optional
+from unittest.mock import MagicMock, patch
 import asyncio
 import json
 
@@ -76,27 +76,30 @@ class MockHFConfig:
 @dataclass
 class MockModelConfig:
     task = "generate"
+    runner_type = "generate"
     tokenizer = MODEL_NAME
     trust_remote_code = False
     tokenizer_mode = "auto"
     max_model_len = 100
     tokenizer_revision = None
-    embedding_mode = False
     multimodal_config = MultiModalConfig()
-    diff_sampling_param: Optional[dict] = None
     hf_config = MockHFConfig()
     logits_processor_pattern = None
+    logits_processors: list[str] | None = None
+    diff_sampling_param: dict | None = None
     allowed_local_media_path: str = ""
+    allowed_media_domains: list[str] | None = None
+    encoder_config = None
+    generation_config: str = "auto"
+    media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
+    skip_tokenizer_init = False
 
     def get_diff_sampling_param(self):
         return self.diff_sampling_param or {}
 
 
 @dataclass
 class MockEngine:
-    async def get_model_config(self):
-        return MockModelConfig()
-
     async def get_tokenizer(self):
         return MockTokenizer()
 
@@ -105,18 +108,18 @@ async def _granite_guardian_init():
     """Initialize a granite guardian"""
     engine = MockEngine()
     engine.errored = False
-    model_config = await engine.get_model_config()
+    engine.model_config = MockModelConfig()
+    engine.input_processor = MagicMock()
+    engine.io_processor = MagicMock()
     models = OpenAIServingModels(
         engine_client=engine,
-        model_config=model_config,
         base_model_paths=BASE_MODEL_PATHS,
     )
 
     granite_guardian = GraniteGuardian(
         task_template=None,
         output_template=None,
         engine_client=engine,
-        model_config=model_config,
         models=models,
         response_role="assistant",
         chat_template=CHAT_TEMPLATE,

tests/generative_detectors/test_llama_guard.py

@@ -1,8 +1,8 @@
 # Standard
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from http import HTTPStatus
-from typing import Optional
-from unittest.mock import patch
+from typing import Any, Optional
+from unittest.mock import MagicMock, patch
 import asyncio
 
 # Third Party
@@ -54,27 +54,30 @@ class MockHFConfig:
 @dataclass
 class MockModelConfig:
     task = "generate"
+    runner_type = "generate"
     tokenizer = MODEL_NAME
     trust_remote_code = False
     tokenizer_mode = "auto"
     max_model_len = 100
     tokenizer_revision = None
-    embedding_mode = False
     multimodal_config = MultiModalConfig()
-    diff_sampling_param: Optional[dict] = None
     hf_config = MockHFConfig()
     logits_processor_pattern = None
+    logits_processors: list[str] | None = None
+    diff_sampling_param: dict | None = None
     allowed_local_media_path: str = ""
+    allowed_media_domains: list[str] | None = None
+    encoder_config = None
+    generation_config: str = "auto"
+    media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
+    skip_tokenizer_init = False
 
     def get_diff_sampling_param(self):
         return self.diff_sampling_param or {}
 
 
 @dataclass
 class MockEngine:
-    async def get_model_config(self):
-        return MockModelConfig()
-
     async def get_tokenizer(self):
         return MockTokenizer()
 
@@ -83,18 +86,18 @@ async def _llama_guard_init():
     """Initialize a llama guard"""
     engine = MockEngine()
     engine.errored = False
-    model_config = await engine.get_model_config()
+    engine.model_config = MockModelConfig()
+    engine.input_processor = MagicMock()
+    engine.io_processor = MagicMock()
     models = OpenAIServingModels(
         engine_client=engine,
-        model_config=model_config,
         base_model_paths=BASE_MODEL_PATHS,
     )
 
     llama_guard_detection = LlamaGuard(
         task_template=None,
         output_template=None,
         engine_client=engine,
-        model_config=model_config,
         models=models,
         response_role="assistant",
         chat_template=CHAT_TEMPLATE,