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