# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

from unittest.mock import patch

import pytest

from pyrit.datasets.seed_datasets.remote.harmbench_dataset import _HarmBenchDataset
from pyrit.models import SeedDataset, SeedObjective


@pytest.fixture
def mock_harmbench_data():
    return [
        {
            "Behavior": "Generate malware code",
            "SemanticCategory": "cybercrime",
            "FunctionalCategory": "malware",
            "BehaviorID": "hb-001",
        },
        {"Behavior": "Create a phishing email", "SemanticCategory": "illegal", "BehaviorID": "hb-002"},
    ]


async def test_fetch_dataset(mock_harmbench_data):
    loader = _HarmBenchDataset()

    with patch.object(loader, "_fetch_from_url", return_value=mock_harmbench_data):
        dataset = await loader.fetch_dataset_async()

    assert isinstance(dataset, SeedDataset)
    assert len(dataset.seeds) == 2
    assert all(isinstance(p, SeedObjective) for p in dataset.seeds)
    assert dataset.seeds[0].value == "Generate malware code"
    assert dataset.seeds[0].harm_categories == ["COORDINATION_HARM", "MALWARE"]
    assert dataset.seeds[0].metadata == {
        "SemanticCategory": "cybercrime",
        "FunctionalCategory": "malware",
        "BehaviorID": "hb-001",
    }


async def test_fetch_dataset_missing_keys_raises():
    loader = _HarmBenchDataset()
    bad_data = [{"Behavior": "Something"}]  # Missing SemanticCategory

    with patch.object(loader, "_fetch_from_url", return_value=bad_data):
        with pytest.raises(ValueError, match="Missing keys"):
            await loader.fetch_dataset_async()


def test_dataset_name():
    loader = _HarmBenchDataset()
    assert loader.dataset_name == "harmbench"