Fix tests broken in prior commit

Signed-off-by: Jared O'Connell <[email protected]>

Author: jaredoconnell

tests/unit/data/deserializers/test_synthetic.py

@@ -11,11 +11,14 @@
 import yaml
 from datasets import IterableDataset
 
-from guidellm.data.deserializers.deserializer import DataNotSupportedError
+from guidellm.data import config as config_module
 from guidellm.data.deserializers.synthetic import (
     SyntheticTextDataset,
-    SyntheticTextDatasetConfig,
     SyntheticTextDatasetDeserializer,
+)
+from guidellm.data.schemas import (
+    DataNotSupportedError,
+    SyntheticTextDatasetConfig,
     SyntheticTextPrefixBucketConfig,
 )
 
@@ -409,13 +412,14 @@ def test_load_config_file_yaml(self):
             yaml_path = f.name
 
         try:
-            deserializer = SyntheticTextDatasetDeserializer()
-            config = deserializer._load_config_file(yaml_path)
+            loaded_config = config_module._load_config_file(
+                yaml_path, SyntheticTextDatasetConfig,
+            )
 
-            assert config.prompt_tokens == 60
-            assert config.output_tokens == 15
-            assert config.source == "yaml_test.txt"
-            assert config.prefix_buckets[0].prefix_tokens == 3  # type: ignore [index]
+            assert loaded_config.prompt_tokens == 60
+            assert loaded_config.output_tokens == 15
+            assert loaded_config.source == "yaml_test.txt"
+            assert loaded_config.prefix_buckets[0].prefix_tokens == 3  # type: ignore [index]
         finally:
             Path(yaml_path).unlink()
 
@@ -438,12 +442,13 @@ def test_load_config_file_config_extension(self):
             config_path = f.name
 
         try:
-            deserializer = SyntheticTextDatasetDeserializer()
-            config = deserializer._load_config_file(config_path)
+            loaded_config = config_module._load_config_file(
+                config_path, SyntheticTextDatasetConfig,
+            )
 
-            assert config.prompt_tokens == 90
-            assert config.output_tokens == 35
-            assert config.prefix_buckets[0].prefix_tokens == 2  # type: ignore [index]
+            assert loaded_config.prompt_tokens == 90
+            assert loaded_config.output_tokens == 35
+            assert loaded_config.prefix_buckets[0].prefix_tokens == 2  # type: ignore [index]
         finally:
             Path(config_path).unlink()
 
@@ -454,11 +459,10 @@ def test_load_config_str_json(self):
         ### WRITTEN BY AI ###
         """
         json_str = '{"prompt_tokens": 50, "output_tokens": 25}'
-        deserializer = SyntheticTextDatasetDeserializer()
-        config = deserializer._load_config_str(json_str)
+        loaded_config = config_module._load_config_str(json_str, SyntheticTextDatasetConfig)
 
-        assert config.prompt_tokens == 50
-        assert config.output_tokens == 25
+        assert loaded_config.prompt_tokens == 50
+        assert loaded_config.output_tokens == 25
 
     @pytest.mark.smoke
     def test_load_config_str_key_value(self):
@@ -467,53 +471,53 @@ def test_load_config_str_key_value(self):
         ### WRITTEN BY AI ###
         """
         kv_str = "prompt_tokens=50,output_tokens=25"
-        deserializer = SyntheticTextDatasetDeserializer()
-        config = deserializer._load_config_str(kv_str)
+        loaded_config = config_module._load_config_str(kv_str, SyntheticTextDatasetConfig)
 
-        assert config.prompt_tokens == 50
-        assert config.output_tokens == 25
+        assert loaded_config.prompt_tokens == 50
+        assert loaded_config.output_tokens == 25
 
     @pytest.mark.sanity
     def test_load_config_str_invalid_format(self):
         """Test loading invalid format raises DataNotSupportedError.
 
         ### WRITTEN BY AI ###
         """
-        deserializer = SyntheticTextDatasetDeserializer()
         with pytest.raises(DataNotSupportedError, match="Unsupported string data"):
-            deserializer._load_config_str("invalid_format_string")
+            config_module._load_config_str(
+                "invalid_format_string", SyntheticTextDatasetConfig,
+            )
 
     @pytest.mark.regression
     def test_load_config_file_non_existent(self):
         """Test loading non-existent file returns None.
 
         ### WRITTEN BY AI ###
         """
-        deserializer = SyntheticTextDatasetDeserializer()
-        config = deserializer._load_config_file("/non/existent/path.config")
-        assert config is None
+        loaded_config = config_module._load_config_file(
+            "/non/existent/path.config", SyntheticTextDatasetConfig,
+        )
+        assert loaded_config is None
 
     @pytest.mark.regression
     def test_load_config_str_non_string(self):
         """Test loading non-string returns None.
 
         ### WRITTEN BY AI ###
         """
-        deserializer = SyntheticTextDatasetDeserializer()
-        config = deserializer._load_config_str(123)
-        assert config is None
+        loaded_config = config_module._load_config_str(123, SyntheticTextDatasetConfig)
+        assert loaded_config is None
 
     @pytest.mark.smoke
     def test_call_with_config_object(self, mock_tokenizer):
         """Test calling deserializer with SyntheticTextDatasetConfig.
 
         ### WRITTEN BY AI ###
         """
-        config = SyntheticTextDatasetConfig(prompt_tokens=50, output_tokens=25)
+        config_input = SyntheticTextDatasetConfig(prompt_tokens=50, output_tokens=25)
         deserializer = SyntheticTextDatasetDeserializer()
 
         result = deserializer(
-            data=config,
+            data=config_input,
             data_kwargs={},
             processor_factory=lambda: mock_tokenizer,
             random_seed=42,