fix to allow to work when using nornir-nautobot ivnetory directly

nautobot · Jun 18, 2024 · a758ae7 · a758ae7
1 parent e33a0f4
commit a758ae7
Showing 1 changed file with 16 additions and 12 deletions.
diff --git a/nornir_nautobot/plugins/tasks/dispatcher/default.py b/nornir_nautobot/plugins/tasks/dispatcher/default.py
@@ -43,22 +43,26 @@ def _get_hostname(cls, task: Task, obj=None) -> str:  # pylint: disable=unused-a
 
     @classmethod
     def _get_tcp_port(cls, obj) -> int:
-        custom_field = obj.cf.get("tcp_port")
-        if isinstance(custom_field, int):
-            return custom_field
-        config_context = obj.get_config_context().get("tcp_port")
-        if isinstance(config_context, int):
-            return config_context
+        if hasattr(obj, "cf"):
+            custom_field = obj.cf.get("tcp_port")
+            if isinstance(custom_field, int):
+                return custom_field
+        if hasattr(obj, "get_config_context"):
+            config_context = obj.get_config_context().get("tcp_port")
+            if isinstance(config_context, int):
+                return config_context
         return cls.tcp_port
 
     @classmethod
     def _get_config_injections(cls, obj) -> list:
-        custom_field = obj.cf.get("config_injections")
-        if isinstance(custom_field, str):
-            return custom_field.split(",")
-        config_context = obj.get_config_context().get("config_injections")
-        if isinstance(config_context, str):
-            return config_context.split(",")
+        if hasattr(obj, "cf"):
+            custom_field = obj.cf.get("config_injections")
+            if isinstance(custom_field, str):
+                return custom_field.split(",")
+        if hasattr(obj, "get_config_context"):
+            config_context = obj.get_config_context().get("config_injections")
+            if isinstance(config_context, str):
+                return config_context.split(",")
         return cls.config_injections
 
     @classmethod