From aa52367f03725d42e343c865c4227e3918e49c1c Mon Sep 17 00:00:00 2001
From: Xavier Bustamante Talavera <xavier@bustawin.com>
Date: Tue, 11 Sep 2018 22:51:13 +0200
Subject: [PATCH] Iterate over SQLA Result / Query instead of .all()

---
 ereuse_devicehub/resources/device/views.py |  3 +--
 ereuse_devicehub/resources/lot/models.py   |  2 +-
 ereuse_devicehub/resources/lot/models.pyi  | 14 ++++++++------
 setup.py                                   |  2 +-
 tests/test_lot.py                          |  4 ++--
 tests/test_snapshot.py                     |  2 +-
 6 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/ereuse_devicehub/resources/device/views.py b/ereuse_devicehub/resources/device/views.py
index 7e6edacd..de66fe95 100644
--- a/ereuse_devicehub/resources/device/views.py
+++ b/ereuse_devicehub/resources/device/views.py
@@ -28,5 +28,4 @@ class DeviceView(View):
 
     def find(self, args: dict):
         """Gets many devices."""
-        devices = Device.query.all()
-        return self.schema.jsonify(devices, many=True)
+        return self.schema.jsonify(Device.query, many=True)
diff --git a/ereuse_devicehub/resources/lot/models.py b/ereuse_devicehub/resources/lot/models.py
index 9da9420c..0da6ff85 100644
--- a/ereuse_devicehub/resources/lot/models.py
+++ b/ereuse_devicehub/resources/lot/models.py
@@ -76,7 +76,7 @@ class Lot(Thing):
     @classmethod
     def roots(cls):
         """Gets the lots that are not under any other lot."""
-        return set(cls.query.join(cls.paths).filter(db.func.nlevel(Path.path) == 1).all())
+        return cls.query.join(cls.paths).filter(db.func.nlevel(Path.path) == 1)
 
     def __repr__(self) -> str:
         return '<Lot {0.name} devices={0.devices!r}>'.format(self)
diff --git a/ereuse_devicehub/resources/lot/models.pyi b/ereuse_devicehub/resources/lot/models.pyi
index 386e8b35..c50b8165 100644
--- a/ereuse_devicehub/resources/lot/models.pyi
+++ b/ereuse_devicehub/resources/lot/models.pyi
@@ -1,15 +1,17 @@
 import uuid
 from datetime import datetime
-from typing import Set, Union
+from typing import Iterable, Set, Union
 from uuid import UUID
 
 from sqlalchemy import Column
-from sqlalchemy.orm import relationship
+from sqlalchemy.orm import Query, relationship
 from sqlalchemy_utils import Ltree
 
 from ereuse_devicehub.resources.device.models import Device
 from ereuse_devicehub.resources.models import Thing
 
+LotQuery = Union[Query, Iterable['Lot']]
+
 
 class Lot(Thing):
     id = ...  # type: Column
@@ -26,18 +28,18 @@ class Lot(Thing):
         self.devices = ...  # type: Set[Device]
         self.paths = ...  # type: Set[Path]
 
-    def add_child(self, child: Union['Lot', uuid.UUID]):
+    def add_child(self, child: Union[Lot, uuid.UUID]):
         pass
 
-    def remove_child(self, child: 'Lot'):
+    def remove_child(self, child: Lot):
         pass
 
     @classmethod
-    def roots(cls):
+    def roots(cls) -> LotQuery:
         pass
 
     @property
-    def children(self) -> Set['Lot']:
+    def children(self) -> LotQuery:
         pass
 
 
diff --git a/setup.py b/setup.py
index f4c13fdb..4659e43d 100644
--- a/setup.py
+++ b/setup.py
@@ -44,7 +44,7 @@ setup(
         'psycopg2-binary',
         'python-stdnum',
         'PyYAML',
-        'teal>=0.2.0a13',
+        'teal>=0.2.0a14',
         'requests',
         'requests-toolbelt',
         'sqlalchemy-utils[password, color, phone]',
diff --git a/tests/test_lot.py b/tests/test_lot.py
index 384bc787..20ddbd6a 100644
--- a/tests/test_lot.py
+++ b/tests/test_lot.py
@@ -179,9 +179,9 @@ def test_lot_roots():
     db.session.add_all(lots)
     db.session.flush()
 
-    assert Lot.roots() == {l1, l2, l3}
+    assert set(Lot.roots()) == {l1, l2, l3}
     l1.add_child(l2)
-    assert Lot.roots() == {l1, l3}
+    assert set(Lot.roots()) == {l1, l3}
 
 
 @pytest.mark.usefixtures(conftest.auth_app_context.__name__)
diff --git a/tests/test_snapshot.py b/tests/test_snapshot.py
index 61694e6e..66160a8e 100644
--- a/tests/test_snapshot.py
+++ b/tests/test_snapshot.py
@@ -239,7 +239,7 @@ def test_snapshot_tag_inner_tag(tag_id: str, user: UserClient, app: Devicehub):
     snapshot_and_check(user, b,
                        event_types=(WorkbenchRate.t, AggregateRate.t, BenchmarkProcessor.t))
     with app.app_context():
-        tag, *_ = Tag.query.all()  # type: Tag
+        tag = Tag.query.one()  # type: Tag
         assert tag.device_id == 1, 'Tag should be linked to the first device'