104 lines
3.3 KiB
Python
104 lines
3.3 KiB
Python
import pytest
|
|
from aarecommon.models.models import (
|
|
BeamlineStateEnum,
|
|
BeamMarkCoeffModel,
|
|
DataCollectionParameters,
|
|
DewarAddress,
|
|
MLBoxType,
|
|
MLOutputModel,
|
|
SampleShortInfo,
|
|
)
|
|
|
|
|
|
def test_sample_short_info_methods():
|
|
info = SampleShortInfo(
|
|
db_id=1,
|
|
puck_name="puck1",
|
|
dewar_name="dew1",
|
|
sample_name="sample1",
|
|
run_number=1,
|
|
aaredb_params=DataCollectionParameters(
|
|
totalangle=180, processingresolution=1.5, cloud=True
|
|
),
|
|
user="group1",
|
|
pin=3,
|
|
location=DewarAddress(segment="A", pos=2),
|
|
)
|
|
|
|
addr = info.tell_address()
|
|
assert addr.puck.segment == "A"
|
|
assert addr.puck.pos == 2
|
|
assert addr.pin == 3
|
|
|
|
assert info.loc_str() == "A2-3"
|
|
assert info.loc_str_sort() == "A2-03"
|
|
assert info.aaredb_params is not None
|
|
assert info.aaredb_params.totalangle == 180
|
|
assert info.aaredb_params.processingresolution == 1.5
|
|
|
|
info_no_loc = info.model_copy(update={"location": None})
|
|
assert info_no_loc.loc_str() == "-"
|
|
assert info_no_loc.loc_str_sort() == ""
|
|
|
|
data = {
|
|
"db_id": 2,
|
|
"puck_name": "puck2",
|
|
"dewar_name": "dew2",
|
|
"sample_name": "sample2",
|
|
"run_number": 2,
|
|
"aaredb_params": {"totalrange": 120, "userresolution": 1.8, "cloud": ""},
|
|
"user": "group2",
|
|
"pin": 4,
|
|
"location": {"segment": "B", "pos": 5},
|
|
}
|
|
info2 = SampleShortInfo.from_dict(data)
|
|
assert info2.db_id == 2
|
|
assert info2.location.segment == "B"
|
|
assert info2.aaredb_params is not None
|
|
assert info2.aaredb_params.totalangle == 120
|
|
assert info2.aaredb_params.processingresolution == 1.8
|
|
assert info2.aaredb_params.cloud is True
|
|
|
|
|
|
def test_beam_mark_coeff_model_apply():
|
|
model = BeamMarkCoeffModel(coeff_x=(1.0, 2.0, 5.0), coeff_y=(3.0, 4.0, 6.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)
|
|
|
|
box_model = model.get_box_model(key)
|
|
assert box_model.conf == 0.8
|
|
|
|
assert model.get_box_tuple(key) == (1, 2, 3, 4)
|
|
assert model.get_box_tuple("NonExistent") is None
|
|
|
|
assert model.get_box_tuple_with_conf(key) == (1, 2, 3, 4, 0.8)
|
|
assert model.get_box_tuple_with_conf("NonExistent") is None
|
|
|
|
tuples = model.get_tuples_for_class(MLBoxType.CRYSTAL)
|
|
assert len(tuples) == 1
|
|
assert tuples[0] == (1, 2, 3, 4)
|
|
|
|
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)
|
|
|
|
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():
|
|
assert BeamlineStateEnum.SampleExchange.display_name() == "Sample exchange"
|
|
assert BeamlineStateEnum.Moving.display_name() == "Moving"
|
|
assert BeamlineStateEnum.display_name(None) == "-"
|