Fix two failing E2E dataset providers (ComicJailbreak timeout, VLGuard 401) (#1863)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: romanlutz

pyrit/datasets/seed_datasets/remote/comic_jailbreak_dataset.py

@@ -105,6 +105,7 @@ def __init__(
         ),
         source_type: Literal["public_url", "file"] = "public_url",
         templates: list[str] | None = None,
+        max_examples: int | None = None,
     ) -> None:
         """
         Initialize the ComicJailbreak dataset loader.
@@ -114,13 +115,17 @@ def __init__(
                 at a pinned commit.
             source_type: The type of source ('public_url' or 'file').
             templates: List of template names to include. If None, all 5 templates are used.
+            max_examples: Maximum number of source goals to render. Each goal produces up to
+                ``len(templates)`` image+text pairs. If None, all goals are rendered. Useful for
+                CI and quick validations where rendering all 300 goals × 5 templates is too slow.
 
         Raises:
             ValueError: If any template name is invalid.
         """
         self.source = source
         self.source_type: Literal["public_url", "file"] = source_type
         self.templates = templates or list(self.TEMPLATE_NAMES)
+        self.max_examples = max_examples
 
         invalid = set(self.templates) - set(self.TEMPLATE_NAMES)
         if invalid:
@@ -166,6 +171,7 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
             template_paths[template_name] = await self._fetch_template_async(template_name)
 
         seeds: list[Seed] = []
+        processed_goals = 0
 
         for row_idx, example in enumerate(examples):
             missing_keys = required_keys - example.keys()
@@ -204,6 +210,10 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
                 )
                 seeds.extend(pair)
 
+            processed_goals += 1
+            if self.max_examples is not None and processed_goals >= self.max_examples:
+                break
+
         logger.info(f"Successfully loaded {len(seeds)} seeds from ComicJailbreak dataset")
         return SeedDataset(seeds=seeds, dataset_name=self.dataset_name)
 

pyrit/datasets/seed_datasets/remote/vlguard_dataset.py

@@ -4,6 +4,7 @@
 import asyncio
 import json
 import logging
+import os
 import uuid
 import zipfile
 from enum import Enum
@@ -108,14 +109,14 @@ def __init__(
             categories (list[VLGuardCategory] | None): List of VLGuard categories to filter by.
                 If None, all categories are included.
             token (str | None): HuggingFace authentication token for accessing the gated dataset.
-                If None, uses the default token from the environment or HuggingFace CLI login.
+                If not provided, reads from the ``HUGGINGFACE_TOKEN`` environment variable.
 
         Raises:
             ValueError: If any of the specified categories are invalid.
         """
         self.subset = subset
         self.categories = categories
-        self.token = token
+        self.token = token if token is not None else os.environ.get("HUGGINGFACE_TOKEN")
         self.source = f"https://huggingface.co/datasets/{_HF_REPO_ID}"
 
         if categories is not None:

tests/end_to_end/test_all_datasets.py

@@ -21,9 +21,14 @@
 
 from pyrit.datasets import SeedDatasetProvider
 from pyrit.datasets.seed_datasets.remote import (
+    _ComicJailbreakDataset,
     _HarmBenchMultimodalDataset,
+    _HiXSTestDataset,
     _PromptIntelDataset,
+    _SGXSTestDataset,
     _SIUODataset,
+    _SorryBenchDataset,
+    _VLGuardDataset,
     _VLSUMultimodalDataset,
 )
 from pyrit.models import SeedDataset
@@ -41,6 +46,20 @@
 # due to rate-limiting, so an empty result is expected in some environments.
 _IMAGE_FETCHING_PROVIDERS: set[type] = {_HarmBenchMultimodalDataset, _SIUODataset, _VLSUMultimodalDataset}
 
+# Providers that produce many seeds and would otherwise exceed _TEST_TIMEOUT.
+# Constructed with max_examples to keep CI fast; full coverage runs are out of scope here.
+_LIMITED_EXAMPLES_PROVIDERS: set[type] = {_ComicJailbreakDataset, _VLSUMultimodalDataset}
+
+# Providers backed by HuggingFace-gated datasets. They require both a HUGGINGFACE_TOKEN
+# and that the token's account has accepted each dataset's terms; skipped when no token
+# is present (e.g. when running E2E locally without secrets).
+_HF_GATED_PROVIDERS: set[type] = {
+    _HiXSTestDataset,
+    _SGXSTestDataset,
+    _SorryBenchDataset,
+    _VLGuardDataset,
+}
+
 
 def get_dataset_providers():
     """Helper to get all registered providers for parameterization."""
@@ -85,12 +104,14 @@ async def test_fetch_dataset(self, name, provider_cls):
         # Skip providers that require credentials not available in CI
         if provider_cls == _PromptIntelDataset and not os.environ.get("PROMPTINTEL_API_KEY"):
             pytest.skip("PROMPTINTEL_API_KEY not set")
+        if provider_cls in _HF_GATED_PROVIDERS and not os.environ.get("HUGGINGFACE_TOKEN"):
+            pytest.skip(f"HUGGINGFACE_TOKEN not set (required for gated dataset used by {name})")
 
         logger.info(f"Testing provider: {name}")
 
         try:
-            # Limit examples for slow multimodal providers that fetch many remote images
-            provider = provider_cls(max_examples=6) if provider_cls == _VLSUMultimodalDataset else provider_cls()
+            # Limit examples for slow providers that would otherwise exceed _TEST_TIMEOUT
+            provider = provider_cls(max_examples=6) if provider_cls in _LIMITED_EXAMPLES_PROVIDERS else provider_cls()
 
             dataset = await _fetch_with_retry(provider)
         except Exception as e:

tests/unit/datasets/test_comic_jailbreak_dataset.py

@@ -169,6 +169,27 @@ async def test_fetch_dataset_empty_goal_skipped(self):
             with pytest.raises(ValueError, match="SeedDataset cannot be empty"):
                 await loader.fetch_dataset_async()
 
+    async def test_fetch_dataset_respects_max_examples(self):
+        """max_examples caps the number of source goals that get rendered."""
+        mock_data = [_make_example(Goal=f"Goal {i}") for i in range(5)]
+        loader = _ComicJailbreakDataset(templates=["article"], max_examples=2)
+
+        with (
+            patch.object(loader, "_fetch_from_url", return_value=mock_data),
+            patch.object(loader, "_fetch_template_async", new_callable=AsyncMock, return_value="/fake/template.png"),
+            patch.object(loader, "_render_comic_async", new_callable=AsyncMock, return_value="/fake/rendered.png"),
+        ):
+            dataset = await loader.fetch_dataset_async(cache=False)
+
+        # 2 goals × 1 template × 3 seeds (objective + image + text) = 6
+        assert len(dataset.seeds) == 6
+        goals = {s.metadata["goal"] for s in dataset.seeds if isinstance(s, SeedPrompt)}
+        assert goals == {"Goal 0", "Goal 1"}
+
+    def test_init_default_max_examples_is_none(self):
+        loader = _ComicJailbreakDataset()
+        assert loader.max_examples is None
+
 
 class TestComicJailbreakTemplates:
     """Tests for the COMIC_JAILBREAK_TEMPLATES constant."""

tests/unit/datasets/test_vlguard_dataset.py

@@ -380,3 +380,27 @@ def mock_hf_download(*, repo_id, filename, repo_type, local_dir, token):
 
         assert metadata == test_metadata
         assert result_dir == cache_dir / "test"
+
+
+class TestVLGuardTokenResolution:
+    """Tests for HuggingFace token resolution on _VLGuardDataset."""
+
+    def test_explicit_token_kwarg_used(self):
+        with patch.dict("os.environ", {}, clear=True):
+            loader = _VLGuardDataset(token="kwarg_token")
+            assert loader.token == "kwarg_token"
+
+    def test_falls_back_to_huggingface_token_env(self):
+        with patch.dict("os.environ", {"HUGGINGFACE_TOKEN": "env_token"}):
+            loader = _VLGuardDataset()
+            assert loader.token == "env_token"
+
+    def test_explicit_kwarg_overrides_env(self):
+        with patch.dict("os.environ", {"HUGGINGFACE_TOKEN": "env_token"}):
+            loader = _VLGuardDataset(token="kwarg_token")
+            assert loader.token == "kwarg_token"
+
+    def test_token_is_none_when_neither_set(self):
+        with patch.dict("os.environ", {}, clear=True):
+            loader = _VLGuardDataset()
+            assert loader.token is None