initial Device model implementation

jgarman · jgarman · commit 0b512be2e43f · 2017-03-23T15:23:30.000-04:00
diff --git a/src/cbapi/defense/models.py b/src/cbapi/defense/models.py
@@ -0,0 +1,14 @@
+from ..models import MutableBaseModel, CreatableModelMixin, NewBaseModel
+
+
+class Device(NewBaseModel):
+    urlobject = "/integrationServices/v3/device"
+    primary_key = "deviceId"
+
+    def __init__(self, cb, model_unique_id, initial_data=None):
+        super(Device, self).__init__(cb, model_unique_id, initial_data)
+
+    def _parse(self, obj):
+        if type(obj) == dict and "deviceInfo" in obj:
+            return obj["deviceInfo"]
+
diff --git a/src/cbapi/defense/rest_api.py b/src/cbapi/defense/rest_api.py
@@ -87,6 +87,7 @@ def and_(self, q):
         return self.where(q)
 
     def _count(self):
+        # TODO: FIX
         args = {'limit': -1}
         if self._query:
             args['q'] = self._query
@@ -99,8 +100,9 @@ def _count(self):
     def _search(self, start=0, rows=0):
         # iterate over total result set, 1000 at a time
         args = {}
-        args['fromRow'] = start
-        args['maxRows'] = self._batch_size
+        if start != 0:
+            args['start'] = start
+        args['rows'] = self._batch_size
 
         current = start
         numrows = 0
@@ -114,18 +116,18 @@ def _search(self, start=0, rows=0):
             query_args = convert_query_params(args)
             result = self._cb.get_object(self._doc_class.urlobject, query_parameters=query_args)
 
-            if len(result) == 0:
-                break
+            self._total_results = result.get("totalResults", 0)
+            self._count_valid = True
 
-            for item in result:
+            for item in result.get("results", []):
                 yield item
                 current += 1
                 numrows += 1
                 if rows and numrows == rows:
                     still_querying = False
                     break
 
-            args['offset'] = current + 1
+            args['start'] = current + 1
 
             if len(result) < self._batch_size:
                 break
