def test_get_all_types(self):
fm = ModelFlame(self.testfile)
etypes = {
'quadrupole', 'bpm', 'drift', 'source', 'rfcavity', 'sbend',
'orbtrim', 'solenoid', 'stripper'
}
self.assertEqual(set(fm.get_all_types()), etypes)
def test_get_index_by_type(self):
fm = ModelFlame(self.testfile)
m = fm.machine
all_types = fm.get_all_types()
for n in range(2, len(all_types)):
etyps = [random.choice(all_types) for _ in range(n)]
e = fm.get_index_by_type(type=etyps)
e0 = {t: m.find(type=t) for t in etyps}
self.assertEqual(e, e0)
def test_find(self):
fm = ModelFlame(self.testfile)
m = fm.machine
all_types = fm.get_all_types()
for i in range(2, 20):
e = fm.find(m.conf(i)['name'])[0]
self.assertEqual(i, e)
for ntype in all_types:
e0 = m.find(type=ntype)
e = fm.find(type=ntype)
self.assertEqual(e, e0)
