added test for objects

Author: Daniel Townsend

aragorm/query/__init__.py

@@ -64,6 +64,8 @@ class Objects(
     table instances are returned, rather than just data.
     """
 
+    _output = Output(as_objects=True)
+
     def __str__(self):
         """
         Need to do this without repeating select ...
@@ -73,11 +75,9 @@ def __str__(self):
             column_names=[]
         )
 
-        for attr in ('_limit', '_where', 'order_by'):
+        for attr in ('_limit', '_where', '_output', 'order_by'):
             setattr(select, attr, getattr(self, attr))
 
-        select._output = Output(as_objects=True)
-
         return select.__str__()
 
 

aragorm/query/base.py

@@ -41,18 +41,22 @@ async def run(self, as_dict=True, credentials=None):
         raw = self.response_handler(raw)
 
         output = getattr(self, '_output', None)
-        if output and type(raw) is list:
-            if output.as_list:
-                if len(raw[0].keys()) != 1:
-                    raise ValueError(
-                        'Each row returned more than on value'
-                    )
-                else:
-                    raw = list(
-                        itertools.chain(*[j.values() for j in raw])
-                    )
-            if output.as_json:
-                raw = json.dumps(raw)
+
+        if output:
+            if output.as_objects:
+                raw = [self.table(**columns) for columns in raw]
+            elif type(raw) is list:
+                if output.as_list:
+                    if len(raw[0].keys()) != 1:
+                        raise ValueError(
+                            'Each row returned more than on value'
+                        )
+                    else:
+                        raw = list(
+                            itertools.chain(*[j.values() for j in raw])
+                        )
+                if output.as_json:
+                    raw = json.dumps(raw)
 
         return raw
 

aragorm/table.py

@@ -282,7 +282,9 @@ def alter(cls) -> Alter:
 
     @classmethod
     def objects(cls) -> Objects:
-        return Objects()
+        return Objects(
+            table=cls
+        )
 
     @classmethod
     def exists(cls) -> Exists:

tests/table/test_objects.py

@@ -0,0 +1,14 @@
+from ..base import DBTestCase
+from ..example_project.tables import Pokemon
+
+
+class TestObjects(DBTestCase):
+
+    def test_get_all(self):
+        self.insert_row()
+
+        response = Pokemon.objects().run_sync()
+
+        self.assertTrue(len(response) == 1)
+        self.assertTrue(isinstance(response[0], Pokemon))
+        self.assertTrue(response[0].name == 'pikachu')