Add methods for get or add/update model (#182)

Author: itdependsnetworks

diffsync/__init__.py

@@ -790,19 +790,41 @@ def get_or_instantiate(
         """
         return self.store.get_or_instantiate(model=model, ids=ids, attrs=attrs)
 
+    def get_or_add_model_instance(self, obj: DiffSyncModel) -> Tuple[DiffSyncModel, bool]:
+        """Attempt to get the object with provided obj identifiers or instantiate obj.
+
+        Args:
+            obj: An obj of the DiffSyncModel to get or add.
+
+        Returns:
+            Provides the existing or new object and whether it was created or not.
+        """
+        return self.store.get_or_add_model_instance(obj=obj)
+
     def update_or_instantiate(self, model: Type[DiffSyncModel], ids: Dict, attrs: Dict) -> Tuple[DiffSyncModel, bool]:
         """Attempt to update an existing object with provided ids/attrs or instantiate it with provided identifiers and attrs.
 
         Args:
-            model (DiffSyncModel): The DiffSyncModel to get or create.
-            ids (Dict): Identifiers for the DiffSyncModel to get or create with.
+            model (DiffSyncModel): The DiffSyncModel to update or create.
+            ids (Dict): Identifiers for the DiffSyncModel to update or create with.
             attrs (Dict): Attributes when creating/updating an object if it doesn't exist. Pass in empty dict, if no specific attrs.
 
         Returns:
             Tuple[DiffSyncModel, bool]: Provides the existing or new object and whether it was created or not.
         """
         return self.store.update_or_instantiate(model=model, ids=ids, attrs=attrs)
 
+    def update_or_add_model_instance(self, obj: DiffSyncModel) -> Tuple[DiffSyncModel, bool]:
+        """Attempt to update an existing object with provided obj ids/attrs or instantiate obj.
+
+        Args:
+            instance: An instance of the DiffSyncModel to update or create.
+
+        Returns:
+            Provides the existing or new object and whether it was created or not.
+        """
+        return self.store.update_or_add_model_instance(obj=obj)
+
     def count(self, model: Union[Text, "DiffSyncModel", Type["DiffSyncModel"], None] = None):
         """Count how many objects of one model type exist in the backend store.
 

diffsync/store/__init__.py

@@ -159,6 +159,24 @@ def get_or_instantiate(
 
         return obj, created
 
+    def get_or_add_model_instance(self, obj: "DiffSyncModel") -> Tuple["DiffSyncModel", bool]:
+        """Attempt to get the object with provided obj identifiers or instantiate obj.
+
+        Args:
+            obj: An obj of the DiffSyncModel to get or add.
+
+        Returns:
+            Provides the existing or new object and whether it was added or not.
+        """
+        model = obj.get_type()
+        ids = obj.get_unique_id()
+
+        try:
+            return self.get(model=model, identifier=ids), False
+        except ObjectNotFound:
+            self.add(obj=obj)
+            return obj, True
+
     def update_or_instantiate(
         self, *, model: Type["DiffSyncModel"], ids: Dict, attrs: Dict
     ) -> Tuple["DiffSyncModel", bool]:
@@ -188,6 +206,33 @@ def update_or_instantiate(
 
         return obj, created
 
+    def update_or_add_model_instance(self, obj: "DiffSyncModel") -> Tuple["DiffSyncModel", bool]:
+        """Attempt to update an existing object with provided ids/attrs or instantiate obj.
+
+        Args:
+            instance: An instance of the DiffSyncModel to update or create.
+
+        Returns:
+            Provides the existing or new object and whether it was added or not.
+        """
+        model = obj.get_type()
+        ids = obj.get_unique_id()
+        attrs = obj.get_attrs()
+
+        added = False
+        try:
+            obj = self.get(model=model, identifier=ids)
+        except ObjectNotFound:
+            # Add the object to the diffsync instance
+            self.add(obj=obj)
+            added = True
+
+        # Update existing obj with attrs
+        for attr, value in attrs.items():
+            setattr(obj, attr, value)
+
+        return obj, added
+
     def _get_object_class_and_model(
         self, model: Union[Text, "DiffSyncModel", Type["DiffSyncModel"]]
     ) -> Tuple[Union["DiffSyncModel", Type["DiffSyncModel"], None], str]:

tests/unit/test_diffsync.py

@@ -162,6 +162,72 @@ def test_diffsync_get_or_instantiate_retrieve_existing_object_wo_attrs(generic_d
     assert obj.description is None
 
 
+def test_diffsync_get_or_add_model_instance_create_non_existent_object(generic_diffsync):
+    generic_diffsync.interface = Interface
+    intf_identifiers = {"device_name": "device1", "name": "eth1"}
+    intf = generic_diffsync.interface(**intf_identifiers)
+
+    # Assert that the object does not currently exist.
+    with pytest.raises(ObjectNotFound):
+        generic_diffsync.get(Interface, intf_identifiers)
+
+    obj, created = generic_diffsync.get_or_add_model_instance(intf)
+    assert created
+    assert obj is generic_diffsync.get(Interface, intf_identifiers)
+    assert obj is generic_diffsync.get("interface", intf_identifiers)
+
+
+def test_diffsync_get_or_add_model_instance_retrieve_existing_object(generic_diffsync):
+    intf_identifiers = {"device_name": "device1", "name": "eth1"}
+    intf = Interface(**intf_identifiers)
+    generic_diffsync.add(intf)
+
+    obj, created = generic_diffsync.get_or_add_model_instance(intf)
+    assert obj is intf
+    assert not created
+
+
+def test_diffsync_get_or_add_model_instance_retrieve_existing_object_w_attrs(generic_diffsync):
+    intf_identifiers = {"device_name": "device1", "name": "eth1"}
+    intf_attrs = {"interface_type": "ethernet"}
+    intf_combine = {**intf_identifiers, **intf_attrs}
+    intf = Interface(**intf_combine)
+    generic_diffsync.add(intf)
+
+    obj, created = generic_diffsync.get_or_add_model_instance(intf)
+    assert obj is intf
+    assert not created
+    assert obj.interface_type == "ethernet"
+    assert obj.description is None
+
+
+def test_diffsync_get_or_add_model_instance_retrieve_create_non_existent_w_attrs(generic_diffsync):
+    generic_diffsync.interface = Interface
+    intf_identifiers = {"device_name": "device1", "name": "eth1"}
+    intf_attrs = {"interface_type": "1000base-t", "description": "Testing"}
+    intf_combine = {**intf_identifiers, **intf_attrs}
+    intf = Interface(**intf_combine)
+
+    obj, created = generic_diffsync.get_or_add_model_instance(intf)
+    assert created
+    assert obj.interface_type == "1000base-t"
+    assert obj.description == "Testing"
+    assert obj is generic_diffsync.get(Interface, intf_identifiers)
+    assert obj is generic_diffsync.get("interface", intf_identifiers)
+
+
+def test_diffsync_get_or_add_model_instance_retrieve_existing_object_wo_attrs(generic_diffsync):
+    intf_identifiers = {"device_name": "device1", "name": "eth1"}
+    intf = Interface(**intf_identifiers)
+    generic_diffsync.add(intf)
+
+    obj, created = generic_diffsync.get_or_add_model_instance(intf)
+    assert obj is intf
+    assert not created
+    assert obj.interface_type == "ethernet"
+    assert obj.description is None
+
+
 def test_diffsync_update_or_instantiate_retrieve_existing_object_w_updated_attrs(generic_diffsync):
     intf_identifiers = {"device_name": "device1", "name": "eth1"}
     intf_attrs = {"interface_type": "1000base-t", "description": "Testing"}
@@ -194,6 +260,42 @@ def test_diffsync_update_or_instantiate_create_object_w_attrs(generic_diffsync):
     assert obj.description == "Testing"
 
 
+def test_diffsync_update_or_add_model_instance_retrieve_existing_object_w_updated_attrs(generic_diffsync):
+    intf_identifiers = {"device_name": "device1", "name": "eth1"}
+    intf_attrs = {"interface_type": "1000base-t", "description": "Testing"}
+    intf_combine = {**intf_identifiers, **intf_attrs}
+    intf = Interface(**intf_combine)
+    generic_diffsync.add(intf)
+
+    obj, created = generic_diffsync.update_or_add_model_instance(intf)
+    assert obj is intf
+    assert not created
+    assert obj.interface_type == "1000base-t"
+    assert obj.description == "Testing"
+
+
+def test_diffsync_update_or_add_model_instance_create_object(generic_diffsync):
+    intf_identifiers = {"device_name": "device1", "name": "eth1"}
+    intf = Interface(**intf_identifiers)
+
+    obj, created = generic_diffsync.update_or_add_model_instance(intf)
+    assert created
+    assert obj.interface_type == "ethernet"
+    assert obj.description is None
+
+
+def test_diffsync_update_or_add_model_instance_create_object_w_attrs(generic_diffsync):
+    intf_identifiers = {"device_name": "device1", "name": "eth1"}
+    intf_attrs = {"interface_type": "1000base-t", "description": "Testing"}
+    intf_combine = {**intf_identifiers, **intf_attrs}
+    intf = Interface(**intf_combine)
+
+    obj, created = generic_diffsync.update_or_add_model_instance(intf)
+    assert created
+    assert obj.interface_type == "1000base-t"
+    assert obj.description == "Testing"
+
+
 def test_diffsync_get_with_generic_model(generic_diffsync, generic_diffsync_model):
     generic_diffsync.add(generic_diffsync_model)
     # The generic_diffsync_model has an empty identifier/unique-id