Use diffsync store methods (#267)

* Use diffsync store methods
networktocode · Jun 9, 2022 · e87781e · e87781e
1 parent 1181e25
commit e87781e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/network_importer/adapters/base.py b/network_importer/adapters/base.py
@@ -1,5 +1,6 @@
 """BaseAdapter for the network importer."""
 from diffsync import DiffSync
+from diffsync.exceptions import ObjectNotFound
 from network_importer.models import Site, Device, Interface, IPAddress, Cable, Vlan, Prefix
 
 
@@ -50,8 +51,10 @@ def get_or_create_vlan(self, vlan, site=None):
         modelname = vlan.get_type()
         uid = vlan.get_unique_id()
 
-        if uid in self._data[modelname]:
-            return self._data[modelname][uid], False
+        try:
+            return self.get(modelname, uid), False
+        except ObjectNotFound:
+            pass
 
         self.add(vlan)
         if site:
@@ -72,8 +75,10 @@ def get_or_add(self, obj):
         modelname = obj.get_type()
         uid = obj.get_unique_id()
 
-        if uid in self._data[modelname]:
-            return self._data[modelname][uid], False
+        try:
+            return self.get(modelname, uid), False
+        except ObjectNotFound:
+            pass
 
         self.add(obj)
 

diff --git a/tests/unit/adapters/network_importer/test_load_batfish_interface.py b/tests/unit/adapters/network_importer/test_load_batfish_interface.py
@@ -154,7 +154,7 @@ def test_load_batfish_interface_intf_ether_sub(network_importer_base, site_sfo,
     assert not intf.is_lag_member
     assert intf.allowed_vlans == ["sfo__201"]
 
-    assert "sfo__201" in adapter._data["vlan"]  # pylint: disable=W0212
+    assert any("sfo__201" == str(obj) for obj in adapter.get_all("vlan"))
 
 
 def test_load_batfish_interface_intf_lag_member(network_importer_base, site_sfo, dev_spine1):