mirror of
https://github.com/jiaxiaojunQAQ/OmniSafeBench-MM.git
synced 2026-07-30 08:18:54 +02:00
test: add test for per-attack image directories and improve error handling
- Add test_generate_test_cases.py to verify each attack uses its own image directory - Improve error messages in unified_registry for failed imports with more context - Add null checks in generate_test_cases.py for attack initialization failures - Reduce excessive max_tokens in judge_utils from 130000 to 10000 for efficiency - Add missing ThreadPoolExecutor import in qr/attack.py - Fix missing newline in pyproject.toml
This commit is contained in:
@@ -0,0 +1,53 @@
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from core.data_formats import PipelineConfig
|
||||
from pipeline.generate_test_cases import TestCaseGenerator as CaseGenerator
|
||||
|
||||
|
||||
def test_run_uses_per_attack_image_dir(tmp_path):
|
||||
config = PipelineConfig(
|
||||
system={"output_dir": str(tmp_path)},
|
||||
test_case_generation={
|
||||
"input": {"behaviors_file": "unused_in_test.json"},
|
||||
"attacks": ["jood", "hades"],
|
||||
"attack_params": {"jood": {}, "hades": {}},
|
||||
},
|
||||
batch_size=1,
|
||||
)
|
||||
generator = CaseGenerator(config)
|
||||
|
||||
behaviors = [{"case_id": 1, "original_prompt": "p", "image_path": "img.jpg"}]
|
||||
captured_image_dirs = {}
|
||||
|
||||
def fake_generate_filename(stage_name, **context):
|
||||
attack_name = context["attack_name"]
|
||||
image_dir = tmp_path / attack_name / "images"
|
||||
output_file = tmp_path / attack_name / "test_cases.jsonl"
|
||||
return image_dir, output_file
|
||||
|
||||
def fake_generate_single_attack_test_cases(
|
||||
attack_name,
|
||||
attack_config,
|
||||
behaviors,
|
||||
batch_size,
|
||||
output_file_path,
|
||||
image_save_dir,
|
||||
):
|
||||
captured_image_dirs[attack_name] = Path(image_save_dir)
|
||||
return []
|
||||
|
||||
with patch.object(generator, "validate_config", return_value=True), patch.object(
|
||||
generator, "load_behaviors", return_value=behaviors
|
||||
), patch.object(generator, "_calculate_expected_test_cases", return_value=1), patch.object(
|
||||
generator, "load_results", return_value=[]
|
||||
), patch.object(generator, "_generate_filename", side_effect=fake_generate_filename), patch.object(
|
||||
generator,
|
||||
"generate_single_attack_test_cases",
|
||||
side_effect=fake_generate_single_attack_test_cases,
|
||||
):
|
||||
generator.run()
|
||||
|
||||
assert captured_image_dirs["jood"] == tmp_path / "jood" / "images"
|
||||
assert captured_image_dirs["hades"] == tmp_path / "hades" / "images"
|
||||
assert captured_image_dirs["jood"] != captured_image_dirs["hades"]
|
||||
Reference in New Issue
Block a user