from unittest import mock

import pytest

from addams_bec.bec_ipython_client.plugins.energy_optimizer.addams_energy_optimizer import (
    EnergyOptimizer,
)


@pytest.fixture
def optimizer():
    dm = mock.MagicMock()
    yield EnergyOptimizer(dm)


@pytest.mark.parametrize(
    "start_energy, target_energy, expected",
    [
        (5, 10, [(10, "Si")]),
        (5, 20, [(9, "Rh"), (20, "Rh")]),
        (5, 40, [(9, "Rh"), (22, "Pt"), (40, "Pt")]),
        (5, 5, [(5, "Si")]),
        (5, 4, ValueError),
        (5, 41, ValueError),
        (2, 8, ValueError),
        (18, 40, [(22, "Pt"), (40, "Pt")]),
        (18, 5, [(9, "Si"), (5, "Si")]),
        (18, 10, [(10, "Rh")]),
        (18, 20, [(20, "Rh")]),
        (25, 7, [(22, "Rh"), (9, "Si"), (7, "Si")]),
    ],
)
def test_get_transition_steps(optimizer, start_energy, target_energy, expected):
    if expected == ValueError:
        with pytest.raises(ValueError):
            optimizer.get_transition_steps(start_energy, target_energy)
    else:
        assert optimizer.get_transition_steps(start_energy, target_energy) == expected