From e7d23f7e5ca4f2bfd64c22a6e31261ab96cf8c17 Mon Sep 17 00:00:00 2001 From: itdependsnetworks Date: Fri, 4 Nov 2022 22:15:07 -0400 Subject: [PATCH] Add methods for get or add/update model --- diffsync/__init__.py | 26 ++++++++- diffsync/store/__init__.py | 45 ++++++++++++++++ tests/unit/test_diffsync.py | 102 ++++++++++++++++++++++++++++++++++++ 3 files changed, 171 insertions(+), 2 deletions(-) diff --git a/diffsync/__init__.py b/diffsync/__init__.py index 9a0dd039..ce7cd975 100644 --- a/diffsync/__init__.py +++ b/diffsync/__init__.py @@ -790,12 +790,23 @@ 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: @@ -803,6 +814,17 @@ def update_or_instantiate(self, model: Type[DiffSyncModel], ids: Dict, attrs: Di """ 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. diff --git a/diffsync/store/__init__.py b/diffsync/store/__init__.py index d2de62f9..aed32ede 100644 --- a/diffsync/store/__init__.py +++ b/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]: diff --git a/tests/unit/test_diffsync.py b/tests/unit/test_diffsync.py index 6d2588b5..d8471f3d 100644 --- a/tests/unit/test_diffsync.py +++ b/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