def test_add_element_to_family_and_case_insensitive_retrieval():
e = Element(6.0, "QUAD", "dummy")
e.add_to_family("FAM")
# Lowercase only
assert "fam" in e.families
assert "FAM" not in e.families
assert e.is_in_family("fam")
assert e.is_in_family("FAM")
def test_element_properties_are_None_without_lattice():
e = Element(1.2, "SEXT")
assert e.index is None
assert e.s is None
assert e.cell is None
lat = mock.Mock()
lat.cell_length = None
e._lattice = lat
assert e.cell is None
def test_element_properties_with_lattice():
e1 = Element(3.1, "DRFIT", "d1")
e2 = Element(1.3, "DRFIT", "d2")
lat = Lattice("", symmetry=2)
lat.add_element(e1)
lat.add_element(e2)
assert e1.index == 1
assert e2.index == 2
assert e1.s == 0.0
assert e2.s == 3.1
assert e1.cell == 1
assert e2.cell == 2
def test_create_element():
lat = Lattice("")
e = Element(6.0, "BPM", "bpm1", lat)
assert e.length == 6.0
assert e.type_ == "BPM"
assert e.name == "bpm1"
assert e._lattice == lat
def test_set_lattice_reference():
elem1 = Element(1.0, "BPM")
lat1 = Lattice("one")
elem2 = Element(2.0, "BPM", lattice=lat1)
lat2 = Lattice("two")
assert elem1._lattice is None
assert elem2._lattice == lat1
elem1.set_lattice(lat1)
elem2.set_lattice(lat2)
assert elem1._lattice == lat1
assert elem2._lattice == lat2
def test_add_element_to_lattice():
lat1 = Lattice("lat1")
elem = Element(0.5, "DRIFT")
assert lat1._elements == []
assert elem._lattice is None
lat1.add_element(elem)
assert lat1._elements == [elem]
assert elem._lattice == lat1
lat2 = Lattice("lat2")
lat2.add_element(elem)
assert elem._lattice == lat2
def test_element_representation():
elem = Element(0.1, "BPM")
assert str(elem) == "<Element length 0.1 m, families >"
elem.add_to_family("fam1")
assert str(elem) == "<Element length 0.1 m, families fam1>"
elem.name = "bpm1"
assert str(elem) == "<Element 'bpm1', length 0.1 m, families fam1>"
lat = Lattice("")
lat.add_element(elem)
assert str(elem) == ("<Element 'bpm1', index 1, length 0.1 m, families " "fam1>")
lat.symmetry = 2
assert str(elem) == (
"<Element 'bpm1', index 1, length 0.1 m, cell 1, " "families fam1>"
)
elem.name = None
assert str(elem) == ("<Element index 1, length 0.1 m, cell 1, families " "fam1>")
def test_lattice_cell_properties():
lat = Lattice("", 2)
for i in range(5):
lat.add_element(Element(0.5, "DRIFT"))
assert lat.cell_length == 1.25
assert lat.cell_bounds == [1, 4, 5]
def test_get_family_s(simple_lattice):
assert simple_lattice.get_family_s("family") == [0]
element2 = Element(1.0, "family")
element2.add_to_family("family")
simple_lattice.add_element(element2)
assert simple_lattice.get_family_s("family") == [0, 0]
element3 = Element(2.5, "family")
element3.add_to_family("family")
simple_lattice.add_element(element3)
assert simple_lattice.get_family_s("family") == [0, 0, 1.0]
element4 = Element(0.0, "family")
element4.add_to_family("family")
simple_lattice.add_element(element4)
assert simple_lattice.get_family_s("family") == [0, 0, 1.0, 3.5]
def simple_element(x_device, y_device, mock_sim_data_source, unit_uc,
double_uc):
# A unit conversion object that returns the same as the input.
element = Element(0.0, "BPM", "element1")
element.add_to_family("family")
element.set_data_source(DeviceDataSource(), pytac.LIVE)
element.add_device("x", x_device, unit_uc)
element.add_device("y", y_device, double_uc)
element.set_data_source(mock_sim_data_source, pytac.SIM)
return element