FIX Reject empty WMDP category values (microsoft#1497)

Co-authored-by: Roman Lutz <romanlutz13@gmail.com>

Author: 2 people

pyrit/datasets/executors/question_answer/wmdp_dataset.py

@@ -31,7 +31,7 @@ def fetch_wmdp_dataset(category: Optional[str] = None) -> QuestionAnsweringDatas
     """
     # Determine which subset of data to load
     data_categories = None
-    if not category:  # if category is not specified, read in all 3 subsets of data
+    if category is None:  # if category is not specified, read in all 3 subsets of data
         data_categories = ["wmdp-cyber", "wmdp-bio", "wmdp-chem"]
     elif category not in ["cyber", "bio", "chem"]:
         raise ValueError(f"Invalid Parameter: {category}. Expected 'cyber', 'bio', or 'chem'")

tests/unit/datasets/test_wmdp_dataset.py

@@ -0,0 +1,24 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+from unittest.mock import patch
+
+import pytest
+
+from pyrit.datasets.executors.question_answer.wmdp_dataset import fetch_wmdp_dataset
+
+
+class _EmptySplit:
+    def __len__(self) -> int:
+        return 0
+
+
+def test_fetch_wmdp_dataset_rejects_empty_category():
+    with patch(
+        "pyrit.datasets.executors.question_answer.wmdp_dataset.load_dataset",
+        return_value={"test": _EmptySplit()},
+    ) as mock_load_dataset:
+        with pytest.raises(ValueError, match="Invalid Parameter"):
+            fetch_wmdp_dataset(category="")
+
+    mock_load_dataset.assert_not_called()