def test_can_get_objects_by_name(self):
objects = [Mock(_name="AAA"), Mock(_name="BBB"), Mock(_name="CCC")]
self.atom1._aaa, self.atom2._aaa, self.atom3._aaa = objects
structure = AtomStructure(self.atom1, self.atom2, self.atom3)
self.assertEqual(structure._get("aaa"), set(objects))
self.assertEqual(structure._get("aaa", name="AAA"), {objects[0]})
self.assertEqual(structure._get("aaa", name_regex="(AAA)|(BBB)"),
set(objects[:2]))
def test_can_get_objects_by_water_status(self):
objects = [Mock(_name="HOH"), Mock(_name="BBB"), Mock(_name="WAT")]
self.atom1._ligand, self.atom2._ligand, self.atom3._ligand = objects
structure = AtomStructure(self.atom1, self.atom2, self.atom3)
self.assertEqual(structure._get("ligand"), set(objects))
self.assertEqual(structure._get("ligand", water=False), {objects[1]})
self.atom1._aaa, self.atom2._aaa, self.atom3._aaa = objects
self.assertEqual(structure._get("aaa", water=False), set(objects))
def test_can_get_objects(self):
structure = AtomStructure(self.atom1, self.atom2, self.atom3)
self.atom1._aaa = 10
self.atom2._aaa = 11
self.atom3._aaa = None
self.assertEqual(structure._get("aaa"), {10, 11})
self.atom3._aaa = 10
self.assertEqual(structure._get("aaa"), {10, 11})