fix creation of new policies

jgarman · jgarman · commit 87ec2387c7d1 · 2017-08-07T17:25:56.000-04:00
diff --git a/src/cbapi/defense/models.py b/src/cbapi/defense/models.py
@@ -11,6 +11,13 @@
 
 class DefenseMutableModel(MutableBaseModel):
     _change_object_http_method = "PATCH"
+    _change_object_key_name = None
+
+    def __init__(self, cb, model_unique_id=None, initial_data=None, force_init=False, full_doc=False):
+        super(DefenseMutableModel, self).__init__(cb, model_unique_id=model_unique_id, initial_data=initial_data,
+                                                  force_init=force_init, full_doc=full_doc)
+        if not self._change_object_key_name:
+            self._change_object_key_name = self.primary_key
 
     def _parse(self, obj):
         if type(obj) == dict and self.info_key in obj:
@@ -79,6 +86,12 @@ def _refresh_if_needed(self, request_ret):
                             self._info = message.get(self.info_key)
                             self._full_init = True
                             refresh_required = False
+                        else:
+                            if self._change_object_key_name in message.keys():
+                                # if all we got back was an ID, try refreshing to get the entire record.
+                                log.debug("Only received an ID back from the server, forcing a refresh")
+                                self._info[self.primary_key] = message[self._change_object_key_name]
+                                refresh_required = True
                     else:
                         # "success" is False
                         raise ServerError(request_ret.status_code, message.get("message", ""),
@@ -131,5 +144,6 @@ class Policy(DefenseMutableModel, CreatableModelMixin):
     info_key = "policyInfo"
     swagger_meta_file = "defense/models/policyInfo.yaml"
     _change_object_http_method = "PUT"
+    _change_object_key_name = "policyId"
 
 
diff --git a/src/cbapi/models.py b/src/cbapi/models.py
@@ -116,7 +116,7 @@ def __get__(self, instance, instance_type=None):
 class ObjectFieldDescriptor(FieldDescriptor):
     def __get__(self, instance, instance_type=None):
         ret = super(ObjectFieldDescriptor, self).__get__(instance, instance_type)
-        return json.loads(ret) or {}
+        return ret or {}
 
 
 class DateTimeFieldDescriptor(FieldDescriptor):
