diff --git a/nornir_nautobot/plugins/tasks/dispatcher/default.py b/nornir_nautobot/plugins/tasks/dispatcher/default.py index c090268..5b1f17d 100644 --- 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