Added pytest mark to UTs

Author: TomerG711

tests/unit/preprocess/test_dataset.py

@@ -33,7 +33,7 @@ def tokenizer_mock():
     )
     return tokenizer
 
-
+@pytest.mark.smoke
 @patch(f"{process_dataset.__module__}.guidellm_load_dataset")
 @patch(f"{process_dataset.__module__}.check_load_processor")
 @patch(f"{process_dataset.__module__}.Dataset")
@@ -68,51 +68,51 @@ def test_strategy_handler_called(
         mock_load_dataset.assert_called_once()
         mock_check_processor.assert_called_once()
 
-
+@pytest.mark.sanity
 def test_handle_ignore_strategy_too_short(tokenizer_mock):
     result = handle_ignore_strategy("short", 10, tokenizer_mock)
     assert result is None
     tokenizer_mock.encode.assert_called_with("short")
 
-
+@pytest.mark.sanity
 def test_handle_ignore_strategy_sufficient_length(tokenizer_mock):
     result = handle_ignore_strategy("long prompt", 5, tokenizer_mock)
     assert result == "long prompt"
     tokenizer_mock.encode.assert_called_with("long prompt")
 
-
+@pytest.mark.sanity
 def test_handle_concatenate_strategy_enough_prompts(tokenizer_mock):
     dataset_iter = iter([{"prompt": "longer"}])
     result = handle_concatenate_strategy(
         "short", 10, dataset_iter, "prompt", tokenizer_mock, "\n"
     )
     assert result == "short\nlonger"
 
-
+@pytest.mark.sanity
 def test_handle_concatenate_strategy_not_enough_prompts(tokenizer_mock):
     dataset_iter: Iterator = iter([])
     result = handle_concatenate_strategy(
         "short", 10, dataset_iter, "prompt", tokenizer_mock, ""
     )
     assert result is None
 
-
+@pytest.mark.sanity
 def test_handle_pad_strategy(tokenizer_mock):
     result = handle_pad_strategy("short", 10, tokenizer_mock, "p")
     assert result == "shortppppp"
 
-
+@pytest.mark.sanity
 def test_handle_error_strategy_valid_prompt(tokenizer_mock):
     result = handle_error_strategy("valid prompt", 5, tokenizer_mock)
     assert result == "valid prompt"
     tokenizer_mock.encode.assert_called_with("valid prompt")
 
-
+@pytest.mark.sanity
 def test_handle_error_strategy_too_short_prompt(tokenizer_mock):
     with pytest.raises(PromptTooShortError):
         handle_error_strategy("short", 10, tokenizer_mock)
 
-
+@pytest.mark.smoke
 @patch("guidellm.preprocess.dataset.save_dataset_to_file")
 @patch("guidellm.preprocess.dataset.Dataset")
 @patch("guidellm.preprocess.dataset.guidellm_load_dataset")
@@ -159,7 +159,7 @@ def test_process_dataset_non_empty(
         assert "output_tokens_count" in item
         assert len(tokenizer_mock.encode(item["prompt"])) <= 3
 
-
+@pytest.mark.sanity
 @patch(f"{process_dataset.__module__}.Dataset")
 @patch(f"{process_dataset.__module__}.guidellm_load_dataset")
 @patch(f"{process_dataset.__module__}.check_load_processor")
@@ -188,7 +188,7 @@ def test_process_dataset_empty_after_processing(
     mock_check_processor.assert_called_once()
     mock_dataset_class.from_list.assert_not_called()
 
-
+@pytest.mark.smoke
 @patch(f"{process_dataset.__module__}.push_dataset_to_hub")
 @patch(f"{process_dataset.__module__}.Dataset")
 @patch(f"{process_dataset.__module__}.guidellm_load_dataset")
@@ -221,7 +221,7 @@ def test_process_dataset_push_to_hub_called(
     )
     mock_push.assert_called_once_with("id123", mock_dataset_obj)
 
-
+@pytest.mark.sanity
 @patch(f"{process_dataset.__module__}.push_dataset_to_hub")
 @patch(f"{process_dataset.__module__}.Dataset")
 @patch(f"{process_dataset.__module__}.guidellm_load_dataset")
@@ -253,29 +253,29 @@ def test_process_dataset_push_to_hub_not_called(
     )
     mock_push.assert_not_called()
 
-
+@pytest.mark.regression
 def test_push_dataset_to_hub_success():
     os.environ["HF_TOKEN"] = "token"
     mock_dataset = MagicMock(spec=Dataset)
     push_dataset_to_hub("dataset_id", mock_dataset)
     mock_dataset.push_to_hub.assert_called_once_with("dataset_id", token="token")
 
-
+@pytest.mark.regression
 def test_push_dataset_to_hub_error_no_env():
     if "HF_TOKEN" in os.environ:
         del os.environ["HF_TOKEN"]
     mock_dataset = MagicMock(spec=Dataset)
     with pytest.raises(ValueError, match="hub_dataset_id and HF_TOKEN"):
         push_dataset_to_hub("dataset_id", mock_dataset)
 
-
+@pytest.mark.regression
 def test_push_dataset_to_hub_error_no_id():
     os.environ["HF_TOKEN"] = "token"
     mock_dataset = MagicMock(spec=Dataset)
     with pytest.raises(ValueError, match="hub_dataset_id and HF_TOKEN"):
         push_dataset_to_hub(None, mock_dataset)
 
-
+@pytest.mark.regression
 @patch.object(Path, "mkdir")
 def test_save_dataset_to_file_csv(mock_mkdir):
     mock_dataset = MagicMock(spec=Dataset)
@@ -284,7 +284,7 @@ def test_save_dataset_to_file_csv(mock_mkdir):
     mock_dataset.to_csv.assert_called_once_with(output_path)
     mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
 
-
+@pytest.mark.regression
 @patch.object(Path, "mkdir")
 def test_save_dataset_to_file_csv_capitalized(mock_mkdir):
     mock_dataset = MagicMock(spec=Dataset)
@@ -293,7 +293,7 @@ def test_save_dataset_to_file_csv_capitalized(mock_mkdir):
     mock_dataset.to_csv.assert_called_once_with(output_path)
     mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
 
-
+@pytest.mark.regression
 @patch.object(Path, "mkdir")
 def test_save_dataset_to_file_json(mock_mkdir):
     mock_dataset = MagicMock(spec=Dataset)
@@ -302,7 +302,7 @@ def test_save_dataset_to_file_json(mock_mkdir):
     mock_dataset.to_json.assert_called_once_with(output_path)
     mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
 
-
+@pytest.mark.regression
 @patch.object(Path, "mkdir")
 def test_save_dataset_to_file_json_capitalized(mock_mkdir):
     mock_dataset = MagicMock(spec=Dataset)
@@ -311,7 +311,7 @@ def test_save_dataset_to_file_json_capitalized(mock_mkdir):
     mock_dataset.to_json.assert_called_once_with(output_path)
     mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
 
-
+@pytest.mark.regression
 @patch.object(Path, "mkdir")
 def test_save_dataset_to_file_jsonl(mock_mkdir):
     mock_dataset = MagicMock(spec=Dataset)
@@ -320,7 +320,7 @@ def test_save_dataset_to_file_jsonl(mock_mkdir):
     mock_dataset.to_json.assert_called_once_with(output_path)
     mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
 
-
+@pytest.mark.regression
 @patch.object(Path, "mkdir")
 def test_save_dataset_to_file_jsonl_capitalized(mock_mkdir):
     mock_dataset = MagicMock(spec=Dataset)
@@ -329,7 +329,7 @@ def test_save_dataset_to_file_jsonl_capitalized(mock_mkdir):
     mock_dataset.to_json.assert_called_once_with(output_path)
     mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
 
-
+@pytest.mark.regression
 @patch.object(Path, "mkdir")
 def test_save_dataset_to_file_parquet(mock_mkdir):
     mock_dataset = MagicMock(spec=Dataset)
@@ -338,7 +338,7 @@ def test_save_dataset_to_file_parquet(mock_mkdir):
     mock_dataset.to_parquet.assert_called_once_with(output_path)
     mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
 
-
+@pytest.mark.regression
 @patch.object(Path, "mkdir")
 def test_save_dataset_to_file_unsupported_type(mock_mkdir):
     mock_dataset = MagicMock(spec=Dataset)