39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
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
|