add multi_entity_update_modes support to mockgun

shotgun_api3/lib/mockgun/mockgun.py

@@ -335,7 +335,12 @@ def batch(self, requests):
  results.append(self.create(request["entity_type"], request["data"]))
  elif request["request_type"] == "update":
  # note: Shotgun.update returns a list of a single item
- results.append(self.update(request["entity_type"], request["entity_id"], request["data"])[0])
+ results.append(
+ self.update(request["entity_type"],
+ request["entity_id"],
+ request["data"],
+ request.get("multi_entity_update_modes"))[0]
+ )
  elif request["request_type"] == "delete":
  results.append(self.delete(request["entity_type"], request["entity_id"]))
  else:
@@ -387,13 +392,13 @@ def create(self, entity_type, data, return_fields=None):
 
  return result
 
- def update(self, entity_type, entity_id, data):
+ def update(self, entity_type, entity_id, data, multi_entity_update_modes=None):
  self._validate_entity_type(entity_type)
  self._validate_entity_data(entity_type, data)
  self._validate_entity_exists(entity_type, entity_id)
 
  row = self._db[entity_type][entity_id]
- self._update_row(entity_type, row, data)
+ self._update_row(entity_type, row, data, multi_entity_update_modes)
 
  return [dict((field, item) for field, item in row.items() if field in data or field in ("type", "id"))]
 
@@ -818,13 +823,26 @@ def _row_matches_filters(self, entity_type, row, filters, filter_operator, retir
  else:
  raise ShotgunError("%s is not a valid filter operator" % filter_operator)
 
- def _update_row(self, entity_type, row, data):
+ def _update_row(self, entity_type, row, data, multi_entity_update_modes=None):
  for field in data:
  field_type = self._get_field_type(entity_type, field)
  if field_type == "entity" and data[field]:
  row[field] = {"type": data[field]["type"], "id": data[field]["id"]}
  elif field_type == "multi_entity":
- row[field] = [{"type": item["type"], "id": item["id"]} for item in data[field]]
+ update_mode = multi_entity_update_modes.get(field, "set") if multi_entity_update_modes else "set"
+
+ if update_mode == "add":
+ row[field] += [{"type": item["type"], "id": item["id"]} for item in data[field]]
+ elif update_mode == "remove":
+ row[field] = [
+ item
+ for item in row[field]
+ for new_item in data[field]
+ if item["id"] != new_item["id"]
+ or item["type"] != new_item["type"]
+ ]
+ elif update_mode == "set":
+ row[field] = [{"type": item["type"], "id": item["id"]} for item in data[field]]
  else:
  row[field] = data[field]
 

tests/test_mockgun.py

@@ -270,6 +270,106 @@ def test_find_with_none(self):
  for item in items:
  self.assertTrue(len(item["users"]) > 0)
 
+
+class TestMultiEntityFieldUpdate(unittest.TestCase):
+ """
+ Ensures multi entity field update modes work.
+ """
+
+ def setUp(self):
+ """
+ Creates test data.
+ """
+
+ self._mockgun = Mockgun("https://test.shotgunstudio.com", login="user", password="1234")
+
+ # Create two versions to assign to the shot.
+ self._version1 = self._mockgun.create("Version", {"code": "version1"})
+ self._version2 = self._mockgun.create("Version", {"code": "version2"})
+ self._version3 = self._mockgun.create("Version", {"code": "version3"})
+
+ # remove 'code' field for later comparisons
+ del self._version1["code"]
+ del self._version2["code"]
+ del self._version3["code"]
+
+ # Create playlists
+ self._add_playlist = self._mockgun.create(
+ "Playlist",
+ {"code": "playlist1", "versions": [self._version1, self._version2]}
+ )
+ self._remove_playlist = self._mockgun.create(
+ "Playlist",
+ {"code": "playlist1", "versions": [self._version1, self._version2, self._version3]}
+ )
+ self._set_playlist = self._mockgun.create(
+ "Playlist",
+ {"code": "playlist1", "versions": [self._version1, self._version2]}
+ )
+
+ def test_update_add(self):
+ """
+ Ensures that "add" multi_entity_update_mode works.
+ """
+ self._mockgun.update(
+ "Playlist", self._add_playlist["id"], {"versions": [self._version3]},
+ multi_entity_update_modes={"versions": "add"}
+ )
+
+ playlist = self._mockgun.find_one(
+ "Playlist", [["id", "is", self._add_playlist["id"]]], ["versions"]
+ )
+ self.assertEqual(
+ playlist["versions"], [self._version1, self._version2, self._version3]
+ )
+
+ def test_update_remove(self):
+ """
+ Ensures that "remove" multi_entity_update_mode works.
+ """
+ self._mockgun.update(
+ "Playlist", self._remove_playlist["id"], {"versions": [self._version2]},
+ multi_entity_update_modes={"versions": "remove"}
+ )
+
+ playlist = self._mockgun.find_one(
+ "Playlist", [["id", "is", self._remove_playlist["id"]]], ["versions"]
+ )
+ self.assertEqual(playlist["versions"], [self._version1, self._version3])
+
+ def test_update_set(self):
+ """
+ Ensures that "set" multi_entity_update_mode works.
+ """
+ self._mockgun.update(
+ "Playlist",
+ self._set_playlist["id"],
+ {"versions": [self._version2, self._version3]},
+ multi_entity_update_modes={"versions": "set"}
+ )
+
+ playlist = self._mockgun.find_one(
+ "Playlist", [["id", "is", self._set_playlist["id"]]], ["versions"]
+ )
+ self.assertEqual(playlist["versions"], [self._version2, self._version3])
+
+ def test_batch_update(self):
+ self._mockgun.batch(
+ [
+ {
+ "request_type": "update",
+ "entity_type": "Playlist",
+ "entity_id": self._set_playlist["id"],
+ "data": {"versions": [self._version1, self._version2]},
+ "multi_entity_update_modes": {"versions": "set"}
+ }
+ ]
+ )
+ playlist = self._mockgun.find_one(
+ "Playlist", [["id", "is", self._set_playlist["id"]]], ["versions"]
+ )
+ self.assertEqual(playlist["versions"], [self._version1, self._version2])
+
 
 class TestFilterOperator(unittest.TestCase):
  """