mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-06-24 22:29:56 +02:00
Alexander Myasoedov
42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from agentic_security.probe_data.image_generator import (
|
|
generate_image,
|
|
generate_image_dataset,
|
|
)
|
|
from agentic_security.probe_data.models import ImageProbeDataset, ProbeDataset
|
|
|
|
|
|
@pytest.mark.parametrize("variant", [0, 1, 2, 3])
|
|
def test_generate_image(variant):
|
|
prompt = "Test prompt"
|
|
image_bytes = generate_image(prompt, variant)
|
|
|
|
assert isinstance(image_bytes, bytes)
|
|
assert len(image_bytes) > 0
|
|
|
|
|
|
@patch("agentic_security.probe_data.image_generator.generate_image")
|
|
def test_generate_image_dataset(mock_generate_image):
|
|
mock_generate_image.return_value = b"dummy_image_bytes"
|
|
|
|
prompt = "Test prompt"
|
|
test_dataset_name = "test_dataset"
|
|
test_datasets = [
|
|
ProbeDataset(
|
|
dataset_name=test_dataset_name,
|
|
prompts=[prompt],
|
|
metadata={},
|
|
tokens=[],
|
|
approx_cost=0.0,
|
|
)
|
|
]
|
|
image_datasets = generate_image_dataset(test_datasets)
|
|
|
|
assert len(image_datasets) == 1
|
|
assert isinstance(image_datasets[0], ImageProbeDataset)
|
|
assert image_datasets[0].test_dataset.dataset_name == test_dataset_name
|
|
assert image_datasets[0].image_prompts[0] == b"dummy_image_bytes"
|