101 lines
3.4 KiB
Python
101 lines
3.4 KiB
Python
import pytest
|
|
from aare.common.models import SampleShortInfo, DewarAddress, BeamMarkCoeffModel, MLOutputModel, MLBoxType, BeamlineStateEnum
|
|
|
|
def test_sample_short_info_methods():
|
|
info = SampleShortInfo(
|
|
db_id=1,
|
|
puck_name="puck1",
|
|
dewar_name="dew1",
|
|
sample_name="sample1",
|
|
run_number=1,
|
|
user="group1",
|
|
pin=3,
|
|
location=DewarAddress(segment="A", pos=2),
|
|
)
|
|
|
|
# Test tell_address (line 380)
|
|
addr = info.tell_address()
|
|
assert addr.puck.segment == "A"
|
|
assert addr.puck.pos == 2
|
|
assert addr.pin == 3
|
|
|
|
# Test loc_str (lines 383-386)
|
|
assert info.loc_str() == "A2-3"
|
|
|
|
# Test loc_str_sort (lines 389-392)
|
|
assert info.loc_str_sort() == "A2-03"
|
|
|
|
info_no_loc = info.model_copy(update={"location": None})
|
|
assert info_no_loc.loc_str() == "-"
|
|
assert info_no_loc.loc_str_sort() == ""
|
|
|
|
# Test from_dict (line 396)
|
|
data = {
|
|
"db_id": 2,
|
|
"puck_name": "puck2",
|
|
"dewar_name": "dew2",
|
|
"sample_name": "sample2",
|
|
"run_number": 2,
|
|
"user": "group2",
|
|
"pin": 4,
|
|
"location": {"segment": "B", "pos": 5}
|
|
}
|
|
info2 = SampleShortInfo.from_dict(data)
|
|
assert info2.db_id == 2
|
|
assert info2.location.segment == "B"
|
|
|
|
def test_beam_mark_coeff_model_apply():
|
|
# Test apply (line 414)
|
|
model = BeamMarkCoeffModel(
|
|
coeff_x=(1.0, 2.0, 5.0),
|
|
coeff_y=(3.0, 4.0, 6.0)
|
|
)
|
|
# zoom = 10
|
|
# x = 1.0 * 100 + 2.0 * 10 + 5.0 = 125.0
|
|
# y = 3.0 * 100 + 4.0 * 10 + 6.0 = 346.0
|
|
res = model.apply(10.0)
|
|
assert res.x == 125.0
|
|
assert res.y == 346.0
|
|
|
|
def test_ml_output_model_extra_methods():
|
|
model = MLOutputModel()
|
|
key = model.add_box(MLBoxType.CRYSTAL, (1, 2, 3, 4), 0.8)
|
|
|
|
# Test get_box_model (line 499)
|
|
box_model = model.get_box_model(key)
|
|
assert box_model.conf == 0.8
|
|
|
|
# Test get_box_tuple (lines 502-505)
|
|
assert model.get_box_tuple(key) == (1, 2, 3, 4)
|
|
assert model.get_box_tuple("NonExistent") is None
|
|
|
|
# Test get_box_tuple_with_conf (lines 508-512)
|
|
assert model.get_box_tuple_with_conf(key) == (1, 2, 3, 4, 0.8)
|
|
assert model.get_box_tuple_with_conf("NonExistent") is None
|
|
|
|
# Test get_tuples_for_class (lines 523-527)
|
|
tuples = model.get_tuples_for_class(MLBoxType.CRYSTAL)
|
|
assert len(tuples) == 1
|
|
assert tuples[0] == (1, 2, 3, 4)
|
|
|
|
# Test get_tuples_with_conf_for_class (lines 530-534)
|
|
tuples_conf = model.get_tuples_with_conf_for_class(MLBoxType.CRYSTAL)
|
|
assert len(tuples_conf) == 1
|
|
assert tuples_conf[0] == (1, 2, 3, 4, 0.8)
|
|
|
|
# Test get_class_str (lines 470, 475-481)
|
|
assert MLOutputModel.get_class_str(MLBoxType.LOOP_ALL) == "Loop_all"
|
|
assert MLOutputModel.get_class_str(MLBoxType.PIN) == "Pin"
|
|
assert MLOutputModel.get_class_str(MLBoxType.CRYSTAL) == "Crystal"
|
|
assert MLOutputModel.get_class_str(MLBoxType.LOOP_FACE) == "Loop_face"
|
|
assert MLOutputModel.get_class_str(MLBoxType.ICE) == "Ice"
|
|
assert MLOutputModel.get_class_str(MLBoxType.NEEDLE) == "Needle"
|
|
assert MLOutputModel.get_class_str(100) == "Unknown"
|
|
|
|
def test_beamline_state_enum_display_name():
|
|
# Test display_name (line 556)
|
|
assert BeamlineStateEnum.SampleExchange.display_name() == "Sample exchange"
|
|
assert BeamlineStateEnum.Moving.display_name() == "Moving"
|
|
# Using value for potentially unknown
|
|
assert BeamlineStateEnum.display_name(None) == "-"
|