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 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
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_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]