base.py unit tests

atomium/base.py

@@ -3,14 +3,31 @@
 import re
 
 def get_object_from_filter(obj, components):
+    """Gets the object whose attributes are actually being queried, which may be
+    a different object if there is a chain.
+
+    :param obj: the intial object.
+    :param list components: the components of the original key.
+    :returns: the relevant object"""
+
     components = components[:]
-    if len(components) > 1:
-        if components[-1] != "regex" and not hasattr(obj, f"__{components[-1]}__"):
-            obj = getattr(obj, components[0])
+    while len(components) > 2:
+        obj = getattr(obj, components.pop(0))
+    if len(components) == 2:
+        if components[-1] != "regex":
+            if not hasattr(obj, f"__{components[-1]}__"):
+                obj = getattr(obj, components[0])
     return obj
 
 
 def get_object_attribute_from_filter(obj, components):
+    """Gets the object's value of some attribute based on a list of key
+    components.
+
+    :param obj: the object with attributes.
+    :param list components: the components of the original key.
+    :returns: the value"""
+
     try:
         return getattr(
          obj, components[-1] if hasattr(obj, components[-1]) else components[-2]
@@ -19,6 +36,14 @@ def get_object_attribute_from_filter(obj, components):
 
 
 def attribute_matches_value(attribute, value, components):
+    """Checks if an attribute value matches a given value. The components given
+    will determine whether an exact match is sought, or whether a more complex
+    criterion is used.
+    
+    :param attribute: the value of an object's attribute.
+    :param value: the value to match against.
+    :param list components: the components of the original key.
+    :rtype: ``bool``"""
 
     if components[-1] == "regex":
         return re.match(value, attribute)
@@ -32,6 +57,9 @@ def filter_objects(objects, key, value):
     """Takes a :py:class:`.StructureSet` of objects, and filters them on object
     properties.
 
+    They key can be an attribute of the object, or a complex double-underscore
+    separated chain of attributes.
+
     :param StructreSet objects: the dictionary of objects - the keys are\
     unimportant.
     :param str key: the attribute to search. This can be an attribute of the\
@@ -112,10 +140,14 @@ def __new__(self, *args, **kwargs):
 
 
 class StructureSet:
-    """A data structure for holding sub-structures. It stores them internally
+    """A data structure for holding structures. It stores them internally
     as a dictionary where they keys are IDs (to allow rapid lookup by ID) and
     the values are all structures with that ID (to allow for duplicate IDs).
 
+    Two structure sets can be added together, but they are immutable - the
+    structures they have when they are made is the structures they will always
+    have.
+
     :param \* args: the structures that will make up the StructureSet."""
 
     def __init__(self, *args):

tests/integration/test_file_structure_reading.py

@@ -274,6 +274,15 @@ def test_structure_processing(self):
          atom1, atom2, atom3, atom4, atom5, atom6, atom7, atom8, atom9, atom10,
          atom11, atom12, atom13, atom14, atom15, atom16, atom17, atom18
         })
+        self.assertEqual(chain1.atoms(het__name="ALA"), {
+         atom1, atom2, atom3, atom4, atom5
+        })
+        self.assertEqual(chain1.atoms(het__name="CYS"), {
+         atom6, atom7, atom8, atom9, atom10, atom11
+        })
+        self.assertEqual(chain1.atoms(het__name__regex="CYS|ALA"), {
+         atom1, atom2, atom3, atom4, atom5, atom6, atom7, atom8, atom9, atom10, atom11
+        })
         self.assertTrue(res1.helix)
         self.assertFalse(res1.strand)
         self.assertTrue(res2.helix)

tests/unit/test_base.py

@@ -0,0 +1,171 @@
+from collections import OrderedDict
+from unittest import TestCase
+from unittest.mock import Mock, patch, PropertyMock, MagicMock
+from atomium.base import *
+
+class ObjectFromFilterTests(TestCase):
+
+    def test_can_get_same_object(self):
+        obj = Mock()
+        obj.__lt__ = MagicMock()
+        obj2 = get_object_from_filter(obj, ["height"])
+        self.assertIs(obj, obj2)
+        obj2 = get_object_from_filter(obj, ["height", "regex"])
+        self.assertIs(obj, obj2)
+        obj2 = get_object_from_filter(obj, ["height", "lt"])
+        self.assertIs(obj, obj2)
+
+
+    def test_can_get_chained_object(self):
+        obj = Mock()
+        obj2 = get_object_from_filter(obj, ["o1", "o2", "o3", "height", "regex"])
+        self.assertIs(obj2, obj.o1.o2.o3)
+        obj2 = get_object_from_filter(obj, ["o1", "o2", "o3", "height"])
+        self.assertIs(obj2, obj.o1.o2.o3)
+
+
+
+class AttributeGettingTests(TestCase):
+
+    def test_can_get_basic_attribute(self):
+        obj = Mock(x=10)
+        self.assertEqual(get_object_attribute_from_filter(obj, ["x"]), 10)
+        self.assertEqual(get_object_attribute_from_filter(obj, ["y", "x"]), 10)
+
+
+    def test_can_get_attribute_from_early_chain(self):
+        obj = Mock(x=10)
+        del obj.regex
+        self.assertEqual(get_object_attribute_from_filter(obj, ["x", "regex"]), 10)
+
+
+    def test_can_get_no_attribute(self):
+        obj = Mock(x=10)
+        del obj.y
+        self.assertIsNone(get_object_attribute_from_filter(obj, ["y"]))
+
+
+
+class AttributeMatchingTests(TestCase):
+
+    def test_exact_match(self):
+        self.assertTrue(attribute_matches_value(10, 10, ["height", "xy"]))
+        self.assertFalse(attribute_matches_value(10, 11, ["height", "xy"]))
+
+
+    def test_regex_match(self):
+        self.assertTrue(attribute_matches_value("jon", "jon|joe", ["name", "regex"]))
+        self.assertFalse(attribute_matches_value("jon", "jon|joe", ["name", "rogox"]))
+
+
+    def test_magic_method_match(self):
+        self.assertTrue(attribute_matches_value(12, 10, ["height", "gt"]))
+        self.assertFalse(attribute_matches_value(10, 10, ["height", "gt"]))
+        self.assertTrue(attribute_matches_value(10, 10, ["height", "gte"]))
+
+
+
+class ObjectFilteringTests(TestCase):
+
+    @patch("atomium.base.get_object_from_filter")
+    @patch("atomium.base.get_object_attribute_from_filter")
+    @patch("atomium.base.attribute_matches_value")
+    @patch("atomium.base.StructureSet")
+    def test_can_filter_objects(self, mock_s, mock_match, mock_getat, mock_getob):
+        structures=[
+         Mock(x="A", y=1), Mock(x="B", y=3), Mock(x="B", y=3),
+         Mock(x="C", y=2), Mock(x="D", y=4), Mock(x="D", y=4)
+        ]
+        objects = Mock(structures=structures)
+        mock_getob.side_effect = lambda s, c: s
+        mock_getat.side_effect = lambda s, c: c[0]
+        mock_match.side_effect = [False, True, False, True, False, False]
+        filter_objects(objects, "key__key2__key_3", "value")
+        for structure in structures:
+            mock_getob.assert_any_call(structure, ["key", "key2", "key_3"])
+            mock_getat.assert_any_call(structure, ["key", "key2", "key_3"])
+            mock_match.assert_any_call("key", "value", ["key", "key2", "key_3"])
+        mock_s.assert_called_with(structures[1], structures[3])
+
+
+
+class QueryDecoratorTests(TestCase):
+
+    def setUp(self):
+        self.s = Mock(structures={2, 4, 6}, ids={1, 3, 5})
+        self.f = lambda s: self.s
+
+
+    def test_can_get_unfiltered_objects(self):
+        f = query(self.f)
+        self.assertEqual(f(self), {2, 4, 6})
+
+
+    @patch("atomium.base.filter_objects")
+    def test_can_get_filtered_objects(self, mock_filter):
+        mock_filter.side_effect = [Mock(structures={20}, ids={10})]
+        f = query(self.f)
+        self.assertEqual(f(self, a=1), {20})
+        mock_filter.assert_any_call(self.s, "a", 1)
+
+
+    @patch("atomium.base.filter_objects")
+    def test_can_get_filtered_objects_as_tuple(self, mock_filter):
+        mock_filter.side_effect = [Mock(structures={2}, ids={1})]
+        f = query(self.f, tuple_=True)
+        self.assertEqual(f(self, a=1), (2,))
+        mock_filter.assert_any_call(self.s, "a", 1)
+
+
+    def test_can_get_objects_by_id(self):
+        f = query(self.f)
+        self.assertEqual(f(self, 3), {self.s.get.return_value})
+        self.s.get.assert_called_with(3)
+        self.assertEqual(f(self, 8), set())
+
+
+
+class GetOneDecoratorTests(TestCase):
+
+    def test_can_get_one(self):
+        f = lambda s: [4, 6, 7]
+        f = getone(f)
+        self.assertEqual(f(self), 4)
+
+
+    def test_can_get_mone(self):
+        f = lambda s: []
+        f = getone(f)
+        self.assertEqual(f(self), None)
+
+
+
+class StructureClassMetaclassTests(TestCase):
+
+    @patch("atomium.base.query")
+    @patch("atomium.base.getone")
+    def test_structure_class_metaclass(self, mock_getone, mock_query):
+        class TestClass(metaclass=StructureClass):
+            def a(self): return 1000
+            def chains(self): return {1: 2, 3: 4}
+            def residues(self): return {10: 2, 30: 4}
+            def ligands(self): return {11: 2, 31: 4}
+            def waters(self): return {12: 2, 32: 4}
+            def molecules(self): return {13: 2, 33: 4}
+            def atoms(self): return {14: 2, 34: 4}
+            def b(self): return 2000
+        obj = TestClass()
+        self.assertIs(obj.chains, mock_query.return_value)
+        self.assertIs(obj.chain, mock_getone.return_value)
+        self.assertIs(obj.residues, mock_query.return_value)
+        self.assertIs(obj.residue, mock_getone.return_value)
+        self.assertIs(obj.ligands, mock_query.return_value)
+        self.assertIs(obj.ligand, mock_getone.return_value)
+        self.assertIs(obj.waters, mock_query.return_value)
+        self.assertIs(obj.water, mock_getone.return_value)
+        self.assertIs(obj.molecules, mock_query.return_value)
+        self.assertIs(obj.molecule, mock_getone.return_value)
+        self.assertIs(obj.atoms, mock_query.return_value)
+        self.assertIs(obj.atom, mock_getone.return_value)
+        self.assertEqual(obj.a(), 1000)
+        self.assertEqual(obj.b(), 2000)

tests/unit/test_structure_sets.py

@@ -0,0 +1,59 @@
+from unittest import TestCase
+from unittest.mock import Mock, patch, PropertyMock, MagicMock
+from atomium.base import StructureSet
+
+class StructureSetTest(TestCase):
+
+    def setUp(self):
+        self.structures = [Mock(_id=1), Mock(_id=2), Mock(_id=2)]
+
+
+
+class StructureSetCreationTests(StructureSetTest):
+
+    def test_can_create_structure_set(self):
+        s = StructureSet(*self.structures)
+        self.assertEqual(s._d, {1: {self.structures[0]}, 2: set(self.structures[1:])})
+
+
+
+class StructureSetAdditionTests(StructureSetTest):
+
+    def test_can_add_structure_sets(self):
+        s1 = StructureSet(*self.structures[:2])
+        s2 = StructureSet(self.structures[2])
+        s = s1 + s2
+        self.assertEqual(s._d, {1: {self.structures[0]}, 2: set(self.structures[1:])})
+
+
+
+class StructureSetLengthTests(StructureSetTest):
+
+    def test_can_get_structure_set_len(self):
+        s = StructureSet(*self.structures)
+        self.assertEqual(len(s), 3)
+
+
+
+class StructureSetIdsTests(StructureSetTest):
+
+    def test_can_get_ids(self):
+        s = StructureSet(*self.structures)
+        self.assertEqual(s.ids, {1, 2})
+
+
+
+class StructureSetStructuresTests(StructureSetTest):
+
+    def test_can_get_structures(self):
+        s = StructureSet(*self.structures)
+        self.assertEqual(set(s.structures), set(self.structures))
+
+
+
+class StructureSetGettingTests(StructureSetTest):
+
+    def test_can_get_structure_by_id(self):
+        s = StructureSet(*self.structures)
+        self.assertEqual(s.get(1), self.structures[0])
+        self.assertIn(s.get(2), self.structures[1:])