add tests for grabbing the dataset

Signed-off-by: dalthecow <dalcowboiz@gmail.com>

Author: DaltheCow

src/guidellm/presentation/data_models.py

@@ -89,9 +89,7 @@ def from_data(cls, request_loader: Any):
                 if creator == SyntheticDatasetCreator:
                     data_dict = SyntheticDatasetConfig.parse_str(data)
                     dataset_name = data_dict.source
-                if creator == FileDatasetCreator or isinstance(
-                    creator, HFDatasetsCreator
-                ):
+                if creator == FileDatasetCreator or creator == HFDatasetsCreator:
                     dataset_name = data
                 if creator == InMemoryDatasetCreator:
                     dataset_name = "In-memory"

tests/unit/presentation/test_data_models.py

@@ -1,6 +1,12 @@
 import pytest
+from unittest.mock import MagicMock, patch
 
-from guidellm.presentation.data_models import Bucket
+from guidellm.dataset.file import FileDatasetCreator
+from guidellm.dataset.hf_datasets import HFDatasetsCreator
+from guidellm.dataset.in_memory import InMemoryDatasetCreator
+from guidellm.dataset.synthetic import SyntheticDatasetCreator
+from guidellm.presentation.data_models import Bucket, Dataset
+from tests.unit.mock_benchmark import mock_generative_benchmark
 
 
 @pytest.mark.smoke
@@ -18,3 +24,80 @@ def test_bucket_from_data():
     assert buckets[1].value == 8.0
     assert buckets[1].count == 5
     assert bucket_width == 1
+
+def mock_processor(cls):
+    return mock_generative_benchmark().request_loader.processor
+
+def new_handle_create(cls, *args, **kwargs):
+    return MagicMock()
+
+def new_extract_dataset_name(cls, *args, **kwargs):
+    return "data:prideandprejudice.txt.gz"
+
+@pytest.mark.smoke
+def test_dataset_from_data_uses_extracted_dataset_name():
+    mock_benchmark = mock_generative_benchmark()
+    with (
+        patch.object(SyntheticDatasetCreator, 'handle_create', new=new_handle_create),
+        patch.object(SyntheticDatasetCreator, 'extract_dataset_name', new=new_extract_dataset_name)
+    ):
+        dataset = Dataset.from_data(mock_benchmark.request_loader)
+        assert dataset.name == "data:prideandprejudice.txt.gz"
+        # with unittest.mock.patch.object(PreTrainedTokenizerBase, 'processor', new=mock_processor):
+
+def new_is_supported(cls, *args, **kwargs):
+    return True
+
+@pytest.mark.smoke
+def test_dataset_from_data_with_in_memory_dataset():
+    mock_benchmark = mock_generative_benchmark()
+    with patch.object(InMemoryDatasetCreator, 'is_supported', new=new_is_supported):
+        dataset = Dataset.from_data(mock_benchmark.request_loader)
+        assert dataset.name == "In-memory"
+
+def hardcoded_isnt_supported(cls, *args, **kwargs):
+    return False
+
+def new_extract_dataset_name_none(cls, *args, **kwargs):
+    return None
+
+@pytest.mark.smoke
+def test_dataset_from_data_with_synthetic_dataset():
+    mock_benchmark = mock_generative_benchmark()
+    with (
+        patch.object(SyntheticDatasetCreator, 'handle_create', new=new_handle_create),
+        patch.object(InMemoryDatasetCreator, 'is_supported', new=hardcoded_isnt_supported),
+        patch.object(SyntheticDatasetCreator, 'is_supported', new=new_is_supported),
+        patch.object(SyntheticDatasetCreator, 'extract_dataset_name', new=new_extract_dataset_name_none)
+    ):
+        dataset = Dataset.from_data(mock_benchmark.request_loader)
+        assert dataset.name == "data:prideandprejudice.txt.gz"
+
+@pytest.mark.smoke
+def test_dataset_from_data_with_file_dataset():
+    mock_benchmark = mock_generative_benchmark()
+    mock_benchmark.request_loader.data = 'dataset.yaml'
+    with (
+        patch.object(FileDatasetCreator, 'handle_create', new=new_handle_create),
+        patch.object(InMemoryDatasetCreator, 'is_supported', new=hardcoded_isnt_supported),
+        patch.object(SyntheticDatasetCreator, 'is_supported', new=hardcoded_isnt_supported),
+        patch.object(FileDatasetCreator, 'is_supported', new=new_is_supported),
+        patch.object(FileDatasetCreator, 'extract_dataset_name', new=new_extract_dataset_name_none)
+    ):
+        dataset = Dataset.from_data(mock_benchmark.request_loader)
+        assert dataset.name == "dataset.yaml"
+
+@pytest.mark.smoke
+def test_dataset_from_data_with_hf_dataset():
+    mock_benchmark = mock_generative_benchmark()
+    mock_benchmark.request_loader.data = 'openai/gsm8k'
+    with (
+        patch.object(HFDatasetsCreator, 'handle_create', new=new_handle_create),
+        patch.object(InMemoryDatasetCreator, 'is_supported', new=hardcoded_isnt_supported),
+        patch.object(SyntheticDatasetCreator, 'is_supported', new=hardcoded_isnt_supported),
+        patch.object(FileDatasetCreator, 'is_supported', new=hardcoded_isnt_supported),
+        patch.object(HFDatasetsCreator, 'is_supported', new=new_is_supported),
+        patch.object(HFDatasetsCreator, 'extract_dataset_name', new=new_extract_dataset_name_none)
+    ):
+        dataset = Dataset.from_data(mock_benchmark.request_loader)
+        assert dataset.name == "openai/gsm8k"