mirror of
https://github.com/jiaxiaojunQAQ/OmniSafeBench-MM.git
synced 2026-04-23 21:36:07 +02:00
Liao, Jie
a4875c66b8
- 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
54 lines
1.9 KiB
Python
54 lines
1.9 KiB
Python
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"]
|