core: models: query tests in Suite class

Now one can get all tests belonging to a Suite. Although that's
probably a huge test set, it shouldn't time out the backend server.

2 files changed, 11 insertions, 8 deletions

diff --git a/squad_client/core/models.py b/squad_client/core/models.py
index 8b3425c..287a75f 100644
--- a/squad_client/core/models.py
+++ b/squad_client/core/models.py
@@ -61,7 +61,8 @@ class SquadObject:
 
         attrs = obj.attrs if len(obj.attrs) else [attr for attr in result.keys()]
         for attr in attrs:
-            setattr(obj, attr.replace(' ', '_').replace('/', '_').replace('-', '_'), result[attr])
+            if attr in result.keys():
+                setattr(obj, attr.replace(' ', '_').replace('/', '_').replace('-', '_'), result[attr])
 
     def __fill__(self, klass, results):
 
@@ -81,7 +82,7 @@ class SquadObject:
 
         return '%s(%s)' % (class_name, ', '.join(attrs_str))
 
-    def __fetch__(self, klass=None, filters=None, count=DEFAULT_COUNT):
+    def __fetch__(self, klass=None, filters=None, count=DEFAULT_COUNT, endpoint=None):
         """
             Generic get method to retrieve objects from API
             count: number of objects to fetch, defaults to 50,
@@ -99,7 +100,7 @@ class SquadObject:
 
         filters['limit'] = count if count < settings.SQUAD_MAX_PAGE_LIMIT else settings.SQUAD_MAX_PAGE_LIMIT
         objects = {}
-        url = klass.endpoint
+        url = endpoint or klass.endpoint
         while url and len(objects) < count:
             response = SquadApi.get(url, filters)
             result = response.json()
@@ -555,15 +556,17 @@ class Suite(SquadObject):
     endpoint = '/api/suites/'
     attrs = ['id', 'slug', 'name', 'project']
 
-    __tests__ = {}
+    __tests__ = None
 
     def add_test(self, test):
         if self.__tests__ is None:
             self.__tests__ = {}
-        self.__tests__[test.id] = test
+        self.__tests__[test.__id__] = test
 
-    @property
-    def tests(self):
+    def tests(self, count=ALL, **filters):
+        if self.__tests__ is None and hasattr(self, 'id') and self.id is not None:
+            endpoint = '%s%d/tests' % (self.endpoint, self.id)
+            self.__tests__ = self.__fetch__(Test, filters, count, endpoint=endpoint)
         return self.__tests__
 
 
diff --git a/squad_client/shortcuts.py b/squad_client/shortcuts.py
index a1879e0..d9a698f 100644
--- a/squad_client/shortcuts.py
+++ b/squad_client/shortcuts.py
@@ -29,7 +29,7 @@ def retrieve_build_results(build_url):
             testrun = testruns[_id]
             test_suites = {}
             for suite in testrun.test_suites:
-                test_suites[suite.name] = {t.short_name: t.status for t in suite.tests.values()}
+                test_suites[suite.name] = {t.short_name: t.status for t in suite.tests().values()}
 
             metric_suites = {}
             for suite in testrun.metric_suites: