From a146c75c38068dcc1b2b3bafcbd86f55d2fbf333 Mon Sep 17 00:00:00 2001
From: Nitin Kumar <nitinkr@juniper.net>
Date: Wed, 20 May 2020 19:56:21 +0530
Subject: [PATCH 1/6] handle file located at any location

---
 salt/modules/junos.py | 638 +++++++++++++++++++++---------------------
 salt/states/junos.py  |   2 +-
 2 files changed, 326 insertions(+), 314 deletions(-)

diff --git a/salt/modules/junos.py b/salt/modules/junos.py
index 9ac637da27e..03568389073 100644
--- a/salt/modules/junos.py
+++ b/salt/modules/junos.py
@@ -33,6 +33,7 @@
 import yaml
 from salt.exceptions import MinionError
 from salt.ext import six
+import tempfile
 
 try:
     from lxml import etree
@@ -72,6 +73,71 @@
 __proxyenabled__ = ["junos"]
 
 
+class HandleFileCopy:
+    """
+    To figure out proper path either from proxy local file system
+    or proxy cache or on master. If required, then only copy from
+    master to proxy
+
+    """
+    def __init__(self, path, **kwargs):
+        self._file_path = path
+        self._cached_folder = None
+        self._cached_file = None
+        self._kwargs = kwargs
+
+    def __enter__(self):
+        if self._file_path.startswith("salt://"):
+            # check if file exists in cache
+            local_cache_path = __salt__["cp.is_cached"](self._file_path)
+            if local_cache_path:
+                master_hash = __salt__["cp.hash_file"](self._file_path)
+                proxy_hash = __salt__["file.get_hash"](local_cache_path)
+                # check if hash is same, else copy newly
+                if master_hash.get("hsum") == proxy_hash:
+                    if self._kwargs:
+                        self._cached_file = salt.utils.files.mkstemp()
+                        with open(self._cached_file, "w") as fp:
+                            template_string = __salt__["slsutil.renderer"](
+                                path=local_cache_path, default_renderer="jinja",
+                                **self._kwargs)
+                            fp.write(template_string)
+                        return self._cached_file
+                    else:
+                        return local_cache_path
+                # continue for else part
+            self._cached_folder = tempfile.mkdtemp()
+            log.debug("Caching file {0} at {1}".format(self._file_path, self._cached_folder))
+            if self._kwargs:
+                self._cached_file = __salt__["cp.get_template"](self._file_path, self._cached_folder,
+                                                          **self._kwargs)
+            else:
+                self._cached_file = __salt__["cp.get_file"](self._file_path, self._cached_folder)
+            if self._cached_file != "":
+                return self._cached_file
+        else:
+            # check for local location of file
+            if __salt__["file.file_exists"](self._file_path):
+                if self._kwargs:
+                    self._cached_file = salt.utils.files.mkstemp()
+                    with open(self._cached_file, "w") as fp:
+                        template_string = __salt__["slsutil.renderer"](
+                            path=self._file_path, default_renderer="jinja",
+                            **self._kwargs)
+                        fp.write(template_string)
+                    return self._cached_file
+                else:
+                    return self._file_path
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        if self._cached_file is not None:
+            salt.utils.files.safe_rm(self._cached_file)
+            log.debug("Deleted cached file: {0}".format(self._cached_file))
+        if self._cached_folder is not None:
+            __salt__["file.rmdir"](self._cached_folder)
+            log.debug("Deleted cached folder: {0}".format(self._cached_folder))
+
+
 def __virtual__():
     """
     We need the Junos adapter libraries for this
@@ -873,156 +939,144 @@ def install_config(path=None, **kwargs):
     if "template_vars" in op:
         kwargs = op["template_vars"]
 
-    try:
-        template_cached_path = salt.utils.files.mkstemp()
-        __salt__["cp.get_template"](path, template_cached_path, **kwargs)
-    except Exception as ex:  # pylint: disable=broad-except
-        ret["message"] = (
-            "Salt failed to render the template, please check file path and syntax."
-            "\nError: {0}".format(str(ex))
-        )
-        ret["out"] = False
-        return ret
-
-    if not os.path.isfile(template_cached_path):
-        ret["message"] = "Invalid file path."
-        ret["out"] = False
-        return ret
-
-    if os.path.getsize(template_cached_path) == 0:
-        ret["message"] = "Template failed to render"
-        ret["out"] = False
-        return ret
+    with HandleFileCopy(path, **kwargs) as template_cached_path:
+        if template_cached_path is None:
+            ret["message"] = "Invalid file path."
+            ret["out"] = False
+            return ret
 
-    write_diff = ""
-    if "diffs_file" in op and op["diffs_file"] is not None:
-        write_diff = op["diffs_file"]
-        del op["diffs_file"]
-
-    op["path"] = template_cached_path
-
-    if "format" not in op:
-        if path.endswith("set"):
-            template_format = "set"
-        elif path.endswith("xml"):
-            template_format = "xml"
-        elif path.endswith("json"):
-            template_format = "json"
-        else:
-            template_format = "text"
+        if os.path.getsize(template_cached_path) == 0:
+            ret["message"] = "Template failed to render"
+            ret["out"] = False
+            return ret
 
-        op["format"] = template_format
+        write_diff = ""
+        if "diffs_file" in op and op["diffs_file"] is not None:
+            write_diff = op["diffs_file"]
+            del op["diffs_file"]
 
-    if "replace" in op and op["replace"]:
-        op["merge"] = False
-        del op["replace"]
-    elif "overwrite" in op and op["overwrite"]:
-        op["overwrite"] = True
-    elif "overwrite" in op and not op["overwrite"]:
-        op["merge"] = True
-        del op["overwrite"]
+        op["path"] = template_cached_path
 
-    db_mode = op.pop("mode", "exclusive")
-    if write_diff and db_mode in ["dynamic", "ephemeral"]:
-        ret[
-            "message"
-        ] = "Write diff is not supported with dynamic/ephemeral configuration mode"
-        ret["out"] = False
-        return ret
+        if "format" not in op:
+            if path.endswith("set"):
+                template_format = "set"
+            elif path.endswith("xml"):
+                template_format = "xml"
+            elif path.endswith("json"):
+                template_format = "json"
+            else:
+                template_format = "text"
 
-    config_params = {}
-    if "ephemeral_instance" in op:
-        config_params["ephemeral_instance"] = op.pop("ephemeral_instance")
-    try:
-        with Config(conn, mode=db_mode, **config_params) as cu:
-            try:
-                cu.load(**op)
-            except Exception as exception:  # pylint: disable=broad-except
-                ret["message"] = 'Could not load configuration due to : "{0}"'.format(
-                    exception
-                )
-                ret["format"] = op["format"]
-                ret["out"] = False
-                return ret
-            finally:
-                salt.utils.files.safe_rm(template_cached_path)
+            op["format"] = template_format
 
-            config_diff = None
-            if db_mode in ["dynamic", "ephemeral"]:
-                log.warning("diff is not supported for dynamic and ephemeral")
-            else:
-                config_diff = cu.diff()
-                if config_diff is None:
-                    ret["message"] = "Configuration already applied!"
-                    ret["out"] = True
-                    return ret
+        if "replace" in op and op["replace"]:
+            op["merge"] = False
+            del op["replace"]
+        elif "overwrite" in op and op["overwrite"]:
+            op["overwrite"] = True
+        elif "overwrite" in op and not op["overwrite"]:
+            op["merge"] = True
+            del op["overwrite"]
 
-            commit_params = {}
-            if "confirm" in op:
-                commit_params["confirm"] = op["confirm"]
-            if "comment" in op:
-                commit_params["comment"] = op["comment"]
+        db_mode = op.pop("mode", "exclusive")
+        if write_diff and db_mode in ["dynamic", "ephemeral"]:
+            ret[
+                "message"
+            ] = "Write diff is not supported with dynamic/ephemeral configuration mode"
+            ret["out"] = False
+            return ret
 
-            # Assume commit_check succeeds and initialize variable check
-            check = True
-            if db_mode in ["dynamic", "ephemeral"]:
-                log.warning("commit check not supported for dynamic and ephemeral")
-            else:
+        config_params = {}
+        if "ephemeral_instance" in op:
+            config_params["ephemeral_instance"] = op.pop("ephemeral_instance")
+        try:
+            with Config(conn, mode=db_mode, **config_params) as cu:
                 try:
-                    check = cu.commit_check()
+                    cu.load(**op)
                 except Exception as exception:  # pylint: disable=broad-except
-                    ret[
-                        "message"
-                    ] = 'Commit check threw the following exception: "{0}"'.format(
+                    ret["message"] = 'Could not load configuration due to : "{0}"'.format(
                         exception
                     )
+                    ret["format"] = op["format"]
                     ret["out"] = False
                     return ret
 
-            if check and not test:
-                try:
-                    cu.commit(**commit_params)
-                    ret["message"] = "Successfully loaded and committed!"
-                except Exception as exception:  # pylint: disable=broad-except
+                config_diff = None
+                if db_mode in ["dynamic", "ephemeral"]:
+                    log.warning("diff is not supported for dynamic and ephemeral")
+                else:
+                    config_diff = cu.diff()
+                    if config_diff is None:
+                        ret["message"] = "Configuration already applied!"
+                        ret["out"] = True
+                        return ret
+
+                commit_params = {}
+                if "confirm" in op:
+                    commit_params["confirm"] = op["confirm"]
+                if "comment" in op:
+                    commit_params["comment"] = op["comment"]
+
+                # Assume commit_check succeeds and initialize variable check
+                check = True
+                if db_mode in ["dynamic", "ephemeral"]:
+                    log.warning("commit check not supported for dynamic and ephemeral")
+                else:
+                    try:
+                        check = cu.commit_check()
+                    except Exception as exception:  # pylint: disable=broad-except
+                        ret[
+                            "message"
+                        ] = 'Commit check threw the following exception: "{0}"'.format(
+                            exception
+                        )
+                        ret["out"] = False
+                        return ret
+
+                if check and not test:
+                    try:
+                        cu.commit(**commit_params)
+                        ret["message"] = "Successfully loaded and committed!"
+                    except Exception as exception:  # pylint: disable=broad-except
+                        ret[
+                            "message"
+                        ] = 'Commit check successful but commit failed with "{0}"'.format(
+                            exception
+                        )
+                        ret["out"] = False
+                        return ret
+                elif not check:
+                    cu.rollback()
                     ret[
                         "message"
-                    ] = 'Commit check successful but commit failed with "{0}"'.format(
+                    ] = "Loaded configuration but commit check failed, hence rolling back configuration."
+                    ret["out"] = False
+                else:
+                    cu.rollback()
+                    ret[
+                        "message"
+                    ] = "Commit check passed, but skipping commit for dry-run and rolling back configuration."
+                    ret["out"] = True
+                try:
+                    if write_diff and config_diff is not None:
+                        with salt.utils.files.fopen(write_diff, "w") as fp:
+                            fp.write(salt.utils.stringutils.to_str(config_diff))
+                except Exception as exception:  # pylint: disable=broad-except
+                    ret["message"] = 'Could not write into diffs_file due to: "{0}"'.format(
                         exception
                     )
                     ret["out"] = False
-                    return ret
-            elif not check:
-                cu.rollback()
-                ret[
-                    "message"
-                ] = "Loaded configuration but commit check failed, hence rolling back configuration."
-                ret["out"] = False
-            else:
-                cu.rollback()
-                ret[
-                    "message"
-                ] = "Commit check passed, but skipping commit for dry-run and rolling back configuration."
-                ret["out"] = True
-            try:
-                if write_diff and config_diff is not None:
-                    with salt.utils.files.fopen(write_diff, "w") as fp:
-                        fp.write(salt.utils.stringutils.to_str(config_diff))
-            except Exception as exception:  # pylint: disable=broad-except
-                ret["message"] = 'Could not write into diffs_file due to: "{0}"'.format(
-                    exception
-                )
-                ret["out"] = False
-    except ValueError as ex:
-        message = "install_config failed due to: {0}".format(str(ex))
-        log.error(message)
-        ret["message"] = message
-        ret["out"] = False
-    except LockError as ex:
-        log.error("Configuration database is locked")
-        ret["message"] = ex.message
-        ret["out"] = False
+        except ValueError as ex:
+            message = "install_config failed due to: {0}".format(str(ex))
+            log.error(message)
+            ret["message"] = message
+            ret["out"] = False
+        except LockError as ex:
+            log.error("Configuration database is locked")
+            ret["message"] = ex.message
+            ret["out"] = False
 
-    return ret
+        return ret
 
 
 def zeroize():
@@ -1150,39 +1204,27 @@ def install_os(path=None, **kwargs):
         ret["out"] = False
         return ret
 
+    install_status = False
     if not no_copy_:
-        # To handle invalid image path scenario
-        try:
-            image_cached_path = salt.utils.files.mkstemp()
-            __salt__["cp.get_file"](path, image_cached_path)
-
-            if not os.path.isfile(image_cached_path):
-                ret["message"] = "Invalid image path."
+        with HandleFileCopy(path) as image_path:
+            if image_path is None:
+                ret["message"] = "Invalid path. Please provide a valid image path"
                 ret["out"] = False
                 return ret
-
-            if os.path.getsize(image_cached_path) == 0:
-                ret["message"] = "Failed to copy image"
+            try:
+                install_status = conn.sw.install(image_path, progress=True, timeout=timeout, **op)
+            except Exception as exception:  # pylint: disable=broad-except
+                ret["message"] = 'Installation failed due to: "{0}"'.format(exception)
                 ret["out"] = False
                 return ret
-            path = image_cached_path
-        except MinionError:
-            ret["message"] = "Invalid path. Please provide a valid image path"
+    else:
+        try:
+            install_status = conn.sw.install(path, progress=True, timeout=timeout, **op)
+        except Exception as exception:  # pylint: disable=broad-except
+            ret["message"] = 'Installation failed due to: "{0}"'.format(exception)
             ret["out"] = False
             return ret
 
-    # install() should not reboot the device, reboot is handled in the next block
-    install_status = False
-    try:
-        install_status = conn.sw.install(path, progress=True, timeout=timeout, **op)
-    except Exception as exception:  # pylint: disable=broad-except
-        ret["message"] = 'Installation failed due to: "{0}"'.format(exception)
-        ret["out"] = False
-        return ret
-    finally:
-        if not no_copy_:
-            salt.utils.files.safe_rm(image_cached_path)
-
     if install_status is True:
         ret["message"] = "Installed the os."
     else:
@@ -1211,7 +1253,7 @@ def install_os(path=None, **kwargs):
     return ret
 
 
-def file_copy(src=None, dest=None):
+def file_copy(src, dest):
     """
     Copies the file from the local device to the junos device
 
@@ -1231,31 +1273,21 @@ def file_copy(src=None, dest=None):
     ret = {}
     ret["out"] = True
 
-    if src is None:
-        ret["message"] = "Please provide the absolute path of the file to be copied."
-        ret["out"] = False
-        return ret
-    if not os.path.isfile(src):
-        ret["message"] = "Invalid source file path"
-        ret["out"] = False
-        return ret
+    with HandleFileCopy(src) as fp:
+        if fp is None:
+            ret["message"] = "Invalid source file path {0}".format(src)
+            ret["out"] = False
+            return ret
 
-    if dest is None:
-        ret[
-            "message"
-        ] = "Please provide the absolute path of the destination where the file is to be copied."
-        ret["out"] = False
+        try:
+            with SCP(conn, progress=True) as scp:
+                scp.put(fp, dest)
+            ret["message"] = "Successfully copied file from {0} to {1}".format(src, dest)
+        except Exception as exception:  # pylint: disable=broad-except
+            ret["message"] = 'Could not copy file : "{0}"'.format(exception)
+            ret["out"] = False
         return ret
 
-    try:
-        with SCP(conn, progress=True) as scp:
-            scp.put(src, dest)
-        ret["message"] = "Successfully copied file from {0} to {1}".format(src, dest)
-    except Exception as exception:  # pylint: disable=broad-except
-        ret["message"] = 'Could not copy file : "{0}"'.format(exception)
-        ret["out"] = False
-    return ret
-
 
 def lock():
     """
@@ -1389,78 +1421,66 @@ def load(path=None, **kwargs):
     if "template_vars" in op:
         kwargs = op["template_vars"]
 
-    try:
-        template_cached_path = salt.utils.files.mkstemp()
-        __salt__["cp.get_template"](path, template_cached_path, **kwargs)
-    except Exception as ex:  # pylint: disable=broad-except
-        ret["message"] = (
-            "Salt failed to render the template, please check file path and syntax."
-            "\nError: {0}".format(str(ex))
-        )
-        ret["out"] = False
-        return ret
+    with HandleFileCopy(path, **kwargs) as template_cached_path:
+        if template_cached_path is None:
+            ret["message"] = "Invalid file path."
+            ret["out"] = False
+            return ret
 
-    if not os.path.isfile(template_cached_path):
-        ret["message"] = "Invalid file path."
-        ret["out"] = False
-        return ret
+        if os.path.getsize(template_cached_path) == 0:
+            ret["message"] = "Template failed to render"
+            ret["out"] = False
+            return ret
 
-    if os.path.getsize(template_cached_path) == 0:
-        ret["message"] = "Template failed to render"
-        ret["out"] = False
-        return ret
+        op["path"] = template_cached_path
 
-    op["path"] = template_cached_path
+        if "format" not in op:
+            if path.endswith("set"):
+                template_format = "set"
+            elif path.endswith("xml"):
+                template_format = "xml"
+            elif path.endswith("json"):
+                template_format = "json"
+            else:
+                template_format = "text"
+
+            op["format"] = template_format
+
+        # Currently, four config_actions are supported: overwrite, replace, update, merge
+        # Allow only one config_action, providing multiple config_action value is not allowed
+        actions = [
+            item
+            for item in ("overwrite", "replace", "update", "merge")
+            if op.get(item, False)
+        ]
+        if len(list(actions)) > 1:
+            ret["message"] = "Only one config_action is allowed. Provided: {0}".format(
+                actions
+            )
+            ret["out"] = False
+            return ret
 
-    if "format" not in op:
-        if path.endswith("set"):
-            template_format = "set"
-        elif path.endswith("xml"):
-            template_format = "xml"
-        elif path.endswith("json"):
-            template_format = "json"
-        else:
-            template_format = "text"
-
-        op["format"] = template_format
-
-    # Currently, four config_actions are supported: overwrite, replace, update, merge
-    # Allow only one config_action, providing multiple config_action value is not allowed
-    actions = [
-        item
-        for item in ("overwrite", "replace", "update", "merge")
-        if op.get(item, False)
-    ]
-    if len(list(actions)) > 1:
-        ret["message"] = "Only one config_action is allowed. Provided: {0}".format(
-            actions
-        )
-        ret["out"] = False
-        return ret
+        if "replace" in op and op["replace"]:
+            op["merge"] = False
+            del op["replace"]
+        elif "overwrite" in op and op["overwrite"]:
+            op["overwrite"] = True
+        elif "merge" in op and op["merge"]:
+            op["merge"] = True
+        elif "overwrite" in op and not op["overwrite"]:
+            op["merge"] = True
+            del op["overwrite"]
 
-    if "replace" in op and op["replace"]:
-        op["merge"] = False
-        del op["replace"]
-    elif "overwrite" in op and op["overwrite"]:
-        op["overwrite"] = True
-    elif "merge" in op and op["merge"]:
-        op["merge"] = True
-    elif "overwrite" in op and not op["overwrite"]:
-        op["merge"] = True
-        del op["overwrite"]
+        try:
+            conn.cu.load(**op)
+            ret["message"] = "Successfully loaded the configuration."
+        except Exception as exception:  # pylint: disable=broad-except
+            ret["message"] = 'Could not load configuration due to : "{0}"'.format(exception)
+            ret["format"] = op["format"]
+            ret["out"] = False
+            return ret
 
-    try:
-        conn.cu.load(**op)
-        ret["message"] = "Successfully loaded the configuration."
-    except Exception as exception:  # pylint: disable=broad-except
-        ret["message"] = 'Could not load configuration due to : "{0}"'.format(exception)
-        ret["format"] = op["format"]
-        ret["out"] = False
         return ret
-    finally:
-        salt.utils.files.safe_rm(template_cached_path)
-
-    return ret
 
 
 def commit_check():
@@ -1532,6 +1552,8 @@ def get_table(
     .. code-block:: bash
 
         salt 'device_name' junos.get_table RouteTable routes.yml
+        salt 'device_name' junos.get_table EthPortTable ethport.yml table_args='{"interface_name": "ge-3/2/2"}'
+        salt 'device_name' junos.get_table EthPortTable ethport.yml salt://tables
     """
 
     conn = __proxy__["junos.conn"]()
@@ -1553,78 +1575,68 @@ def get_table(
     pyez_tables_path = os.path.dirname(os.path.abspath(tables_dir.__file__))
     try:
         if path is not None:
-            file_loc = glob.glob(os.path.join(path, "{0}".format(table_file)))
+            file_path = os.path.join(path, "{0}".format(table_file))
         else:
-            file_loc = glob.glob(
-                os.path.join(pyez_tables_path, "{0}".format(table_file))
-            )
-        if len(file_loc) == 1:
-            file_name = file_loc[0]
-        elif len(file_loc) > 1:
-            ret[
-                "message"
-            ] = "Given table file {0} is located at multiple location".format(
-                table_file
-            )
-            ret["out"] = False
-            return ret
-        elif len(file_loc) == 0:
-            ret["message"] = "Given table file {0} cannot be located".format(table_file)
-            ret["out"] = False
-            return ret
-        try:
-            with salt.utils.files.fopen(file_name) as fp:
-                ret["table"] = yaml.load(fp.read(), Loader=yamlordereddictloader.Loader)
-                globals().update(FactoryLoader().load(ret["table"]))
-        except IOError as err:
-            ret[
-                "message"
-            ] = "Uncaught exception during YAML Load - please report: {0}".format(
-                six.text_type(err)
-            )
-            ret["out"] = False
-            return ret
-        try:
-            data = globals()[table](conn)
-            data.get(**get_kvargs)
-        except KeyError as err:
-            ret[
-                "message"
-            ] = "Uncaught exception during get API call - please report: {0}".format(
-                six.text_type(err)
-            )
-            ret["out"] = False
-            return ret
-        except ConnectClosedError:
-            ret[
-                "message"
-            ] = "Got ConnectClosedError exception. Connection lost with {}".format(conn)
-            ret["out"] = False
-            return ret
-        ret["reply"] = json.loads(data.to_json())
-        if data.__class__.__bases__[0] in [OpTable, CfgTable]:
-            # Sets key value if not present in YAML. To be used by returner
-            if ret["table"][table].get("key") is None:
-                ret["table"][table]["key"] = data.ITEM_NAME_XPATH
-            # If key is provided from salt state file.
-            if key is not None:
-                ret["table"][table]["key"] = data.KEY
-            if table_args is not None:
-                args = copy.copy(data.GET_ARGS)
-                args.update(table_args)
-                ret["table"][table]["args"] = args
-        else:
-            if target is not None:
-                ret["table"][table]["target"] = data.TARGET
-            if key is not None:
-                ret["table"][table]["key"] = data.KEY
-            if key_items is not None:
-                ret["table"][table]["key_items"] = data.KEY_ITEMS
-            if table_args is not None:
-                args = copy.copy(data.CMD_ARGS)
-                args.update(table_args)
-                ret["table"][table]["args"] = args
-                ret["table"][table]["command"] = data.GET_CMD
+            file_path = os.path.join(pyez_tables_path, "{0}".format(table_file))
+
+        with HandleFileCopy(file_path) as file_loc:
+            if file_loc is None:
+                ret["message"] = "Given table file {0} cannot be located".format(table_file)
+                ret["out"] = False
+                return ret
+            try:
+                with salt.utils.files.fopen(file_loc) as fp:
+                    ret["table"] = yaml.load(fp.read(), Loader=yamlordereddictloader.Loader)
+                    globals().update(FactoryLoader().load(ret["table"]))
+            except IOError as err:
+                ret[
+                    "message"
+                ] = "Uncaught exception during YAML Load - please report: {0}".format(
+                    six.text_type(err)
+                )
+                ret["out"] = False
+                return ret
+            try:
+                data = globals()[table](conn)
+                data.get(**get_kvargs)
+            except KeyError as err:
+                ret[
+                    "message"
+                ] = "Uncaught exception during get API call - please report: {0}".format(
+                    six.text_type(err)
+                )
+                ret["out"] = False
+                return ret
+            except ConnectClosedError:
+                ret[
+                    "message"
+                ] = "Got ConnectClosedError exception. Connection lost with {}".format(conn)
+                ret["out"] = False
+                return ret
+            ret["reply"] = json.loads(data.to_json())
+            if data.__class__.__bases__[0] in [OpTable, CfgTable]:
+                # Sets key value if not present in YAML. To be used by returner
+                if ret["table"][table].get("key") is None:
+                    ret["table"][table]["key"] = data.ITEM_NAME_XPATH
+                # If key is provided from salt state file.
+                if key is not None:
+                    ret["table"][table]["key"] = data.KEY
+                if table_args is not None:
+                    args = copy.copy(data.GET_ARGS)
+                    args.update(table_args)
+                    ret["table"][table]["args"] = args
+            else:
+                if target is not None:
+                    ret["table"][table]["target"] = data.TARGET
+                if key is not None:
+                    ret["table"][table]["key"] = data.KEY
+                if key_items is not None:
+                    ret["table"][table]["key_items"] = data.KEY_ITEMS
+                if table_args is not None:
+                    args = copy.copy(data.CMD_ARGS)
+                    args.update(table_args)
+                    ret["table"][table]["args"] = args
+                    ret["table"][table]["command"] = data.GET_CMD
     except ConnectClosedError:
         ret["message"] = (
             "Got ConnectClosedError exception. Connection lost "
diff --git a/salt/states/junos.py b/salt/states/junos.py
index 735c86c0047..f43b2b40a62 100644
--- a/salt/states/junos.py
+++ b/salt/states/junos.py
@@ -195,7 +195,7 @@ def rollback(name, id, **kwargs):
 
 
 @resultdecorator
-def diff(name, d_id, **kwargs):
+def diff(name, d_id=0, **kwargs):
     """
     .. versionchanged:: Sodium
 

From bc3cd8652bee96896b0224a81e33c72886f82263 Mon Sep 17 00:00:00 2001
From: Nitin Kumar <nitinkr@juniper.net>
Date: Wed, 20 May 2020 22:30:20 +0530
Subject: [PATCH 2/6] fix unit tests as per latest code

---
 salt/modules/junos.py            |  66 ++--
 tests/unit/modules/test_junos.py | 631 ++++++++++++++++++-------------
 2 files changed, 416 insertions(+), 281 deletions(-)

diff --git a/salt/modules/junos.py b/salt/modules/junos.py
index 03568389073..f3622e5050e 100644
--- a/salt/modules/junos.py
+++ b/salt/modules/junos.py
@@ -18,7 +18,6 @@
 from __future__ import absolute_import, print_function, unicode_literals
 
 import copy
-import glob
 import json
 import logging
 import os
@@ -31,7 +30,6 @@
 import salt.utils.json
 import salt.utils.stringutils
 import yaml
-from salt.exceptions import MinionError
 from salt.ext import six
 import tempfile
 
@@ -80,6 +78,7 @@ class HandleFileCopy:
     master to proxy
 
     """
+
     def __init__(self, path, **kwargs):
         self._file_path = path
         self._cached_folder = None
@@ -97,22 +96,29 @@ def __enter__(self):
                 if master_hash.get("hsum") == proxy_hash:
                     if self._kwargs:
                         self._cached_file = salt.utils.files.mkstemp()
-                        with open(self._cached_file, "w") as fp:
+                        with salt.utils.files.fopen(self._cached_file, "w") as fp:
                             template_string = __salt__["slsutil.renderer"](
-                                path=local_cache_path, default_renderer="jinja",
-                                **self._kwargs)
+                                path=local_cache_path,
+                                default_renderer="jinja",
+                                **self._kwargs
+                            )
                             fp.write(template_string)
                         return self._cached_file
                     else:
                         return local_cache_path
                 # continue for else part
             self._cached_folder = tempfile.mkdtemp()
-            log.debug("Caching file {0} at {1}".format(self._file_path, self._cached_folder))
+            log.debug(
+                "Caching file {0} at {1}".format(self._file_path, self._cached_folder)
+            )
             if self._kwargs:
-                self._cached_file = __salt__["cp.get_template"](self._file_path, self._cached_folder,
-                                                          **self._kwargs)
+                self._cached_file = __salt__["cp.get_template"](
+                    self._file_path, self._cached_folder, **self._kwargs
+                )
             else:
-                self._cached_file = __salt__["cp.get_file"](self._file_path, self._cached_folder)
+                self._cached_file = __salt__["cp.get_file"](
+                    self._file_path, self._cached_folder
+                )
             if self._cached_file != "":
                 return self._cached_file
         else:
@@ -120,10 +126,12 @@ def __enter__(self):
             if __salt__["file.file_exists"](self._file_path):
                 if self._kwargs:
                     self._cached_file = salt.utils.files.mkstemp()
-                    with open(self._cached_file, "w") as fp:
+                    with salt.utils.files.fopen(self._cached_file, "w") as fp:
                         template_string = __salt__["slsutil.renderer"](
-                            path=self._file_path, default_renderer="jinja",
-                            **self._kwargs)
+                            path=self._file_path,
+                            default_renderer="jinja",
+                            **self._kwargs
+                        )
                         fp.write(template_string)
                     return self._cached_file
                 else:
@@ -994,9 +1002,9 @@ def install_config(path=None, **kwargs):
                 try:
                     cu.load(**op)
                 except Exception as exception:  # pylint: disable=broad-except
-                    ret["message"] = 'Could not load configuration due to : "{0}"'.format(
-                        exception
-                    )
+                    ret[
+                        "message"
+                    ] = 'Could not load configuration due to : "{0}"'.format(exception)
                     ret["format"] = op["format"]
                     ret["out"] = False
                     return ret
@@ -1062,7 +1070,9 @@ def install_config(path=None, **kwargs):
                         with salt.utils.files.fopen(write_diff, "w") as fp:
                             fp.write(salt.utils.stringutils.to_str(config_diff))
                 except Exception as exception:  # pylint: disable=broad-except
-                    ret["message"] = 'Could not write into diffs_file due to: "{0}"'.format(
+                    ret[
+                        "message"
+                    ] = 'Could not write into diffs_file due to: "{0}"'.format(
                         exception
                     )
                     ret["out"] = False
@@ -1212,7 +1222,9 @@ def install_os(path=None, **kwargs):
                 ret["out"] = False
                 return ret
             try:
-                install_status = conn.sw.install(image_path, progress=True, timeout=timeout, **op)
+                install_status = conn.sw.install(
+                    image_path, progress=True, timeout=timeout, **op
+                )
             except Exception as exception:  # pylint: disable=broad-except
                 ret["message"] = 'Installation failed due to: "{0}"'.format(exception)
                 ret["out"] = False
@@ -1282,7 +1294,9 @@ def file_copy(src, dest):
         try:
             with SCP(conn, progress=True) as scp:
                 scp.put(fp, dest)
-            ret["message"] = "Successfully copied file from {0} to {1}".format(src, dest)
+            ret["message"] = "Successfully copied file from {0} to {1}".format(
+                src, dest
+            )
         except Exception as exception:  # pylint: disable=broad-except
             ret["message"] = 'Could not copy file : "{0}"'.format(exception)
             ret["out"] = False
@@ -1475,7 +1489,9 @@ def load(path=None, **kwargs):
             conn.cu.load(**op)
             ret["message"] = "Successfully loaded the configuration."
         except Exception as exception:  # pylint: disable=broad-except
-            ret["message"] = 'Could not load configuration due to : "{0}"'.format(exception)
+            ret["message"] = 'Could not load configuration due to : "{0}"'.format(
+                exception
+            )
             ret["format"] = op["format"]
             ret["out"] = False
             return ret
@@ -1581,12 +1597,16 @@ def get_table(
 
         with HandleFileCopy(file_path) as file_loc:
             if file_loc is None:
-                ret["message"] = "Given table file {0} cannot be located".format(table_file)
+                ret["message"] = "Given table file {0} cannot be located".format(
+                    table_file
+                )
                 ret["out"] = False
                 return ret
             try:
                 with salt.utils.files.fopen(file_loc) as fp:
-                    ret["table"] = yaml.load(fp.read(), Loader=yamlordereddictloader.Loader)
+                    ret["table"] = yaml.load(
+                        fp.read(), Loader=yamlordereddictloader.Loader
+                    )
                     globals().update(FactoryLoader().load(ret["table"]))
             except IOError as err:
                 ret[
@@ -1610,7 +1630,9 @@ def get_table(
             except ConnectClosedError:
                 ret[
                     "message"
-                ] = "Got ConnectClosedError exception. Connection lost with {}".format(conn)
+                ] = "Got ConnectClosedError exception. Connection lost with {}".format(
+                    conn
+                )
                 ret["out"] = False
                 return ret
             ret["reply"] = json.loads(data.to_json())
diff --git a/tests/unit/modules/test_junos.py b/tests/unit/modules/test_junos.py
index c94eaf5e88d..1a937d8e208 100644
--- a/tests/unit/modules/test_junos.py
+++ b/tests/unit/modules/test_junos.py
@@ -13,7 +13,7 @@
 
 # Import test libs
 from tests.support.mixins import LoaderModuleMockMixin, XMLEqualityMixin
-from tests.support.mock import ANY, PropertyMock, call, mock_open, patch
+from tests.support.mock import ANY, PropertyMock, call, mock_open, patch, MagicMock
 from tests.support.unit import TestCase, skipIf
 
 # Import 3rd-party libs
@@ -47,6 +47,10 @@ def setup_loader_modules(self):
                 "__salt__": {
                     "cp.get_template": self.mock_cp,
                     "cp.get_file": self.mock_cp,
+                    "file.file_exists": MagicMock(return_value=True),
+                    "slsutil.renderer": MagicMock(
+                        return_value="set system host-name dummy"
+                    ),
                 },
             }
         }
@@ -1120,234 +1124,301 @@ def test_install_config_without_args(self):
         self.assertEqual(junos.install_config(), ret)
 
     def test_install_config_cp_fails(self):
-        with patch("os.path.isfile") as mock_isfile:
-            mock_isfile.return_value = False
+        with patch.dict(
+            junos.__salt__, {"file.file_exists": MagicMock(return_value=False)}
+        ):
+            ret = dict()
             ret = dict()
             ret["message"] = "Invalid file path."
             ret["out"] = False
             self.assertEqual(junos.install_config("path"), ret)
 
     def test_install_config_file_cp_fails(self):
-        with patch("os.path.isfile") as mock_isfile, patch(
-            "os.path.getsize"
-        ) as mock_getsize:
-            mock_isfile.return_value = True
-            mock_getsize.return_value = 0
+        with patch.dict(
+            junos.__salt__, {"file.file_exists": MagicMock(return_value=False)}
+        ):
+            ret = dict()
             ret = dict()
-            ret["message"] = "Template failed to render"
+            ret["message"] = "Invalid file path."
             ret["out"] = False
             self.assertEqual(junos.install_config("path"), ret)
 
     def test_install_config(self):
-        with patch("jnpr.junos.utils.config.Config.commit") as mock_commit, patch(
-            "jnpr.junos.utils.config.Config.commit_check"
-        ) as mock_commit_check, patch(
-            "jnpr.junos.utils.config.Config.diff"
-        ) as mock_diff, patch(
-            "jnpr.junos.utils.config.Config.load"
-        ) as mock_load, patch(
-            "salt.utils.files.safe_rm"
-        ) as mock_safe_rm, patch(
-            "salt.utils.files.mkstemp"
-        ) as mock_mkstemp, patch(
-            "os.path.isfile"
-        ) as mock_isfile, patch(
-            "os.path.getsize"
-        ) as mock_getsize:
-            mock_isfile.return_value = True
-            mock_getsize.return_value = 10
-            mock_mkstemp.return_value = "test/path/config"
-            mock_diff.return_value = "diff"
-            mock_commit_check.return_value = True
-
-            ret = dict()
-            ret["message"] = "Successfully loaded and committed!"
-            ret["out"] = True
-            self.assertEqual(junos.install_config("actual/path/config.set"), ret)
-            mock_load.assert_called_with(path="test/path/config", format="set")
+        with patch.dict(
+            junos.__salt__,
+            {
+                "cp.is_cached": MagicMock(return_value="test/path/config"),
+                "cp.hash_file": MagicMock(
+                    return_value={"hash_type": "sha256", "hsum": "a386e49c17"}
+                ),
+                "file.get_hash": MagicMock(return_value="a386e49c17"),
+            },
+        ):
+            with patch("jnpr.junos.utils.config.Config.commit") as mock_commit, patch(
+                "jnpr.junos.utils.config.Config.commit_check"
+            ) as mock_commit_check, patch(
+                "jnpr.junos.utils.config.Config.diff"
+            ) as mock_diff, patch(
+                "jnpr.junos.utils.config.Config.load"
+            ) as mock_load, patch(
+                "salt.utils.files.safe_rm"
+            ) as mock_safe_rm, patch(
+                "salt.utils.files.mkstemp"
+            ) as mock_mkstemp, patch(
+                "os.path.isfile"
+            ) as mock_isfile, patch(
+                "os.path.getsize"
+            ) as mock_getsize:
+                mock_isfile.return_value = True
+                mock_getsize.return_value = 10
+                mock_mkstemp.return_value = "test/path/config"
+                mock_diff.return_value = "diff"
+                mock_commit_check.return_value = True
+
+                ret = dict()
+                ret["message"] = "Successfully loaded and committed!"
+                ret["out"] = True
+                self.assertEqual(
+                    junos.install_config("salt://actual/path/config.set"), ret
+                )
+                mock_load.assert_called_with(path="test/path/config", format="set")
 
     def test_install_config_xml_file(self):
-        with patch("jnpr.junos.utils.config.Config.commit") as mock_commit, patch(
-            "jnpr.junos.utils.config.Config.commit_check"
-        ) as mock_commit_check, patch(
-            "jnpr.junos.utils.config.Config.diff"
-        ) as mock_diff, patch(
-            "jnpr.junos.utils.config.Config.load"
-        ) as mock_load, patch(
-            "salt.utils.files.safe_rm"
-        ) as mock_safe_rm, patch(
-            "salt.utils.files.mkstemp"
-        ) as mock_mkstemp, patch(
-            "os.path.isfile"
-        ) as mock_isfile, patch(
-            "os.path.getsize"
-        ) as mock_getsize:
-            mock_isfile.return_value = True
-            mock_getsize.return_value = 10
-            mock_mkstemp.return_value = "test/path/config"
-            mock_diff.return_value = "diff"
-            mock_commit_check.return_value = True
-
-            ret = dict()
-            ret["message"] = "Successfully loaded and committed!"
-            ret["out"] = True
-            self.assertEqual(junos.install_config("actual/path/config.xml"), ret)
-            mock_load.assert_called_with(path="test/path/config", format="xml")
+        with patch.dict(
+            junos.__salt__,
+            {
+                "cp.is_cached": MagicMock(return_value="test/path/config"),
+                "cp.hash_file": MagicMock(
+                    return_value={"hash_type": "sha256", "hsum": "a386e49c17"}
+                ),
+                "file.get_hash": MagicMock(return_value="a386e49c17"),
+            },
+        ):
+            with patch("jnpr.junos.utils.config.Config.commit") as mock_commit, patch(
+                "jnpr.junos.utils.config.Config.commit_check"
+            ) as mock_commit_check, patch(
+                "jnpr.junos.utils.config.Config.diff"
+            ) as mock_diff, patch(
+                "jnpr.junos.utils.config.Config.load"
+            ) as mock_load, patch(
+                "salt.utils.files.safe_rm"
+            ) as mock_safe_rm, patch(
+                "salt.utils.files.mkstemp"
+            ) as mock_mkstemp, patch(
+                "os.path.isfile"
+            ) as mock_isfile, patch(
+                "os.path.getsize"
+            ) as mock_getsize:
+                mock_isfile.return_value = True
+                mock_getsize.return_value = 10
+                mock_mkstemp.return_value = "test/path/config"
+                mock_diff.return_value = "diff"
+                mock_commit_check.return_value = True
+
+                ret = dict()
+                ret["message"] = "Successfully loaded and committed!"
+                ret["out"] = True
+                self.assertEqual(
+                    junos.install_config("salt://actual/path/config.xml"), ret
+                )
+                mock_load.assert_called_with(path="test/path/config", format="xml")
 
     def test_install_config_text_file(self):
-        with patch("jnpr.junos.utils.config.Config.commit") as mock_commit, patch(
-            "jnpr.junos.utils.config.Config.commit_check"
-        ) as mock_commit_check, patch(
-            "jnpr.junos.utils.config.Config.diff"
-        ) as mock_diff, patch(
-            "jnpr.junos.utils.config.Config.load"
-        ) as mock_load, patch(
-            "salt.utils.files.safe_rm"
-        ) as mock_safe_rm, patch(
-            "salt.utils.files.mkstemp"
-        ) as mock_mkstemp, patch(
-            "os.path.isfile"
-        ) as mock_isfile, patch(
-            "os.path.getsize"
-        ) as mock_getsize:
-            mock_isfile.return_value = True
-            mock_getsize.return_value = 10
-            mock_mkstemp.return_value = "test/path/config"
-            mock_diff.return_value = "diff"
-            mock_commit_check.return_value = True
-
-            ret = dict()
-            ret["message"] = "Successfully loaded and committed!"
-            ret["out"] = True
-            self.assertEqual(junos.install_config("actual/path/config"), ret)
-            mock_load.assert_called_with(path="test/path/config", format="text")
+        with patch.dict(
+            junos.__salt__,
+            {
+                "cp.is_cached": MagicMock(return_value="test/path/config"),
+                "cp.hash_file": MagicMock(
+                    return_value={"hash_type": "sha256", "hsum": "a386e49c17"}
+                ),
+                "file.get_hash": MagicMock(return_value="a386e49c17"),
+            },
+        ):
+            with patch("jnpr.junos.utils.config.Config.commit") as mock_commit, patch(
+                "jnpr.junos.utils.config.Config.commit_check"
+            ) as mock_commit_check, patch(
+                "jnpr.junos.utils.config.Config.diff"
+            ) as mock_diff, patch(
+                "jnpr.junos.utils.config.Config.load"
+            ) as mock_load, patch(
+                "salt.utils.files.safe_rm"
+            ) as mock_safe_rm, patch(
+                "salt.utils.files.mkstemp"
+            ) as mock_mkstemp, patch(
+                "os.path.isfile"
+            ) as mock_isfile, patch(
+                "os.path.getsize"
+            ) as mock_getsize:
+                mock_isfile.return_value = True
+                mock_getsize.return_value = 10
+                mock_mkstemp.return_value = "test/path/config"
+                mock_diff.return_value = "diff"
+                mock_commit_check.return_value = True
+
+                ret = dict()
+                ret["message"] = "Successfully loaded and committed!"
+                ret["out"] = True
+                self.assertEqual(junos.install_config("salt://actual/path/config"), ret)
+                mock_load.assert_called_with(path="test/path/config", format="text")
 
     def test_install_config_replace(self):
-        with patch("jnpr.junos.utils.config.Config.commit") as mock_commit, patch(
-            "jnpr.junos.utils.config.Config.commit_check"
-        ) as mock_commit_check, patch(
-            "jnpr.junos.utils.config.Config.diff"
-        ) as mock_diff, patch(
-            "jnpr.junos.utils.config.Config.load"
-        ) as mock_load, patch(
-            "salt.utils.files.safe_rm"
-        ) as mock_safe_rm, patch(
-            "salt.utils.files.mkstemp"
-        ) as mock_mkstemp, patch(
-            "os.path.isfile"
-        ) as mock_isfile, patch(
-            "os.path.getsize"
-        ) as mock_getsize:
-            mock_isfile.return_value = True
-            mock_getsize.return_value = 10
-            mock_mkstemp.return_value = "test/path/config"
-            mock_diff.return_value = "diff"
-            mock_commit_check.return_value = True
-
-            args = {
-                "__pub_user": "root",
-                "__pub_arg": [{"replace": True}],
-                "replace": True,
-                "__pub_fun": "junos.install_config",
-                "__pub_jid": "20170222213858582619",
-                "__pub_tgt": "mac_min",
-                "__pub_tgt_type": "glob",
-                "__pub_ret": "",
-            }
+        with patch.dict(
+            junos.__salt__,
+            {
+                "cp.is_cached": MagicMock(return_value="test/path/config"),
+                "cp.hash_file": MagicMock(
+                    return_value={"hash_type": "sha256", "hsum": "a386e49c17"}
+                ),
+                "file.get_hash": MagicMock(return_value="a386e49c17"),
+            },
+        ):
+            with patch("jnpr.junos.utils.config.Config.commit") as mock_commit, patch(
+                "jnpr.junos.utils.config.Config.commit_check"
+            ) as mock_commit_check, patch(
+                "jnpr.junos.utils.config.Config.diff"
+            ) as mock_diff, patch(
+                "jnpr.junos.utils.config.Config.load"
+            ) as mock_load, patch(
+                "salt.utils.files.safe_rm"
+            ) as mock_safe_rm, patch(
+                "salt.utils.files.mkstemp"
+            ) as mock_mkstemp, patch(
+                "os.path.isfile"
+            ) as mock_isfile, patch(
+                "os.path.getsize"
+            ) as mock_getsize:
+                mock_isfile.return_value = True
+                mock_getsize.return_value = 10
+                mock_mkstemp.return_value = "test/path/config"
+                mock_diff.return_value = "diff"
+                mock_commit_check.return_value = True
+
+                args = {
+                    "__pub_user": "root",
+                    "__pub_arg": [{"replace": True}],
+                    "replace": True,
+                    "__pub_fun": "junos.install_config",
+                    "__pub_jid": "20170222213858582619",
+                    "__pub_tgt": "mac_min",
+                    "__pub_tgt_type": "glob",
+                    "__pub_ret": "",
+                }
 
-            ret = dict()
-            ret["message"] = "Successfully loaded and committed!"
-            ret["out"] = True
-            self.assertEqual(
-                junos.install_config("actual/path/config.set", **args), ret
-            )
-            mock_load.assert_called_with(
-                path="test/path/config", format="set", merge=False
-            )
+                ret = dict()
+                ret["message"] = "Successfully loaded and committed!"
+                ret["out"] = True
+                self.assertEqual(
+                    junos.install_config("salt://actual/path/config.set", **args), ret
+                )
+                mock_load.assert_called_with(
+                    path="test/path/config", format="set", merge=False
+                )
 
     def test_install_config_overwrite(self):
-        with patch("jnpr.junos.utils.config.Config.commit") as mock_commit, patch(
-            "jnpr.junos.utils.config.Config.commit_check"
-        ) as mock_commit_check, patch(
-            "jnpr.junos.utils.config.Config.diff"
-        ) as mock_diff, patch(
-            "jnpr.junos.utils.config.Config.load"
-        ) as mock_load, patch(
-            "salt.utils.files.safe_rm"
-        ) as mock_safe_rm, patch(
-            "salt.utils.files.mkstemp"
-        ) as mock_mkstemp, patch(
-            "os.path.isfile"
-        ) as mock_isfile, patch(
-            "os.path.getsize"
-        ) as mock_getsize:
-            mock_isfile.return_value = True
-            mock_getsize.return_value = 10
-            mock_mkstemp.return_value = "test/path/config"
-            mock_diff.return_value = "diff"
-            mock_commit_check.return_value = True
+        with patch.dict(
+            junos.__salt__,
+            {
+                "cp.is_cached": MagicMock(return_value="test/path/config"),
+                "cp.hash_file": MagicMock(
+                    return_value={"hash_type": "sha256", "hsum": "a386e49c17"}
+                ),
+                "file.get_hash": MagicMock(return_value="a386e49c17"),
+            },
+        ):
+            with patch("jnpr.junos.utils.config.Config.commit") as mock_commit, patch(
+                "jnpr.junos.utils.config.Config.commit_check"
+            ) as mock_commit_check, patch(
+                "jnpr.junos.utils.config.Config.diff"
+            ) as mock_diff, patch(
+                "jnpr.junos.utils.config.Config.load"
+            ) as mock_load, patch(
+                "salt.utils.files.safe_rm"
+            ) as mock_safe_rm, patch(
+                "salt.utils.files.mkstemp"
+            ) as mock_mkstemp, patch(
+                "os.path.isfile"
+            ) as mock_isfile, patch(
+                "os.path.getsize"
+            ) as mock_getsize:
+                mock_isfile.return_value = True
+                mock_getsize.return_value = 10
+                mock_mkstemp.return_value = "test/path/config"
+                mock_diff.return_value = "diff"
+                mock_commit_check.return_value = True
+
+                args = {
+                    "__pub_user": "root",
+                    "__pub_arg": [{"overwrite": True}],
+                    "overwrite": True,
+                    "__pub_fun": "junos.install_config",
+                    "__pub_jid": "20170222213858582619",
+                    "__pub_tgt": "mac_min",
+                    "__pub_tgt_type": "glob",
+                    "__pub_ret": "",
+                }
 
-            args = {
-                "__pub_user": "root",
-                "__pub_arg": [{"overwrite": True}],
-                "overwrite": True,
-                "__pub_fun": "junos.install_config",
-                "__pub_jid": "20170222213858582619",
-                "__pub_tgt": "mac_min",
-                "__pub_tgt_type": "glob",
-                "__pub_ret": "",
-            }
-
-            ret = dict()
-            ret["message"] = "Successfully loaded and committed!"
-            ret["out"] = True
-            self.assertEqual(
-                junos.install_config("actual/path/config.xml", **args), ret
-            )
-            mock_load.assert_called_with(
-                path="test/path/config", format="xml", overwrite=True
-            )
+                ret = dict()
+                ret["message"] = "Successfully loaded and committed!"
+                ret["out"] = True
+                self.assertEqual(
+                    junos.install_config("salt://actual/path/config.xml", **args), ret
+                )
+                mock_load.assert_called_with(
+                    path="test/path/config", format="xml", overwrite=True
+                )
 
     def test_install_config_overwrite_false(self):
-        with patch("jnpr.junos.utils.config.Config.commit") as mock_commit, patch(
-            "jnpr.junos.utils.config.Config.commit_check"
-        ) as mock_commit_check, patch(
-            "jnpr.junos.utils.config.Config.diff"
-        ) as mock_diff, patch(
-            "jnpr.junos.utils.config.Config.load"
-        ) as mock_load, patch(
-            "salt.utils.files.safe_rm"
-        ) as mock_safe_rm, patch(
-            "salt.utils.files.mkstemp"
-        ) as mock_mkstemp, patch(
-            "os.path.isfile"
-        ) as mock_isfile, patch(
-            "os.path.getsize"
-        ) as mock_getsize:
-            mock_isfile.return_value = True
-            mock_getsize.return_value = 10
-            mock_mkstemp.return_value = "test/path/config"
-            mock_diff.return_value = "diff"
-            mock_commit_check.return_value = True
-
-            args = {
-                "__pub_user": "root",
-                "__pub_arg": [{"overwrite": False}],
-                "overwrite": False,
-                "__pub_fun": "junos.install_config",
-                "__pub_jid": "20170222213858582619",
-                "__pub_tgt": "mac_min",
-                "__pub_tgt_type": "glob",
-                "__pub_ret": "",
-            }
+        with patch.dict(
+            junos.__salt__,
+            {
+                "cp.is_cached": MagicMock(return_value="test/path/config"),
+                "cp.hash_file": MagicMock(
+                    return_value={"hash_type": "sha256", "hsum": "a386e49c17"}
+                ),
+                "file.get_hash": MagicMock(return_value="a386e49c17"),
+            },
+        ):
+            with patch("jnpr.junos.utils.config.Config.commit") as mock_commit, patch(
+                "jnpr.junos.utils.config.Config.commit_check"
+            ) as mock_commit_check, patch(
+                "jnpr.junos.utils.config.Config.diff"
+            ) as mock_diff, patch(
+                "jnpr.junos.utils.config.Config.load"
+            ) as mock_load, patch(
+                "salt.utils.files.safe_rm"
+            ) as mock_safe_rm, patch(
+                "salt.utils.files.mkstemp"
+            ) as mock_mkstemp, patch(
+                "os.path.isfile"
+            ) as mock_isfile, patch(
+                "os.path.getsize"
+            ) as mock_getsize:
+                mock_isfile.return_value = True
+                mock_getsize.return_value = 10
+                mock_mkstemp.return_value = "test/path/config"
+                mock_diff.return_value = "diff"
+                mock_commit_check.return_value = True
+
+                args = {
+                    "__pub_user": "root",
+                    "__pub_arg": [{"overwrite": False}],
+                    "overwrite": False,
+                    "__pub_fun": "junos.install_config",
+                    "__pub_jid": "20170222213858582619",
+                    "__pub_tgt": "mac_min",
+                    "__pub_tgt_type": "glob",
+                    "__pub_ret": "",
+                }
 
-            ret = dict()
-            ret["message"] = "Successfully loaded and committed!"
-            ret["out"] = True
-            self.assertEqual(junos.install_config("actual/path/config", **args), ret)
-            mock_load.assert_called_with(
-                path="test/path/config", format="text", merge=True
-            )
+                ret = dict()
+                ret["message"] = "Successfully loaded and committed!"
+                ret["out"] = True
+                self.assertEqual(
+                    junos.install_config("salt://actual/path/config", **args), ret
+                )
+                mock_load.assert_called_with(
+                    path="test/path/config", format="text", merge=True
+                )
 
     def test_install_config_load_causes_exception(self):
         with patch("jnpr.junos.utils.config.Config.diff") as mock_diff, patch(
@@ -1685,24 +1756,33 @@ def test_install_os_without_args(self):
         self.assertEqual(junos.install_os(), ret)
 
     def test_install_os_cp_fails(self):
-        with patch("os.path.isfile") as mock_isfile, patch(
-            "os.path.getsize"
-        ) as mock_getsize:
-            mock_getsize.return_value = 10
-            mock_isfile.return_value = False
-            ret = dict()
-            ret["message"] = "Invalid image path."
-            ret["out"] = False
-            self.assertEqual(junos.install_os("/image/path/"), ret)
+        with patch.dict(
+            junos.__salt__,
+            {
+                "cp.is_cached": MagicMock(return_value="/pat/to/tmp/file"),
+                "cp.hash_file": MagicMock(
+                    return_value={"hash_type": "sha256", "hsum": "a386e49c17"}
+                ),
+                "file.get_hash": MagicMock(return_value="xxxx"),
+                "file.rmdir": MagicMock(return_value="True"),
+            },
+        ):
+            with patch("os.path.isfile") as mock_isfile, patch(
+                "os.path.getsize"
+            ) as mock_getsize:
+                mock_getsize.return_value = 10
+                mock_isfile.return_value = False
+                ret = dict()
+                ret["message"] = "Invalid path. Please provide a valid image path"
+                ret["out"] = False
+                self.assertEqual(junos.install_os("salt://image/path/"), ret)
 
     def test_install_os_image_cp_fails(self):
-        with patch("os.path.isfile") as mock_isfile, patch(
-            "os.path.getsize"
-        ) as mock_getsize:
-            mock_getsize.return_value = 0
-            mock_isfile.return_value = True
+        with patch.dict(
+            junos.__salt__, {"file.file_exists": MagicMock(return_value=False)}
+        ):
             ret = dict()
-            ret["message"] = "Failed to copy image"
+            ret["message"] = "Invalid path. Please provide a valid image path"
             ret["out"] = False
             self.assertEqual(junos.install_os("/image/path/"), ret)
 
@@ -1882,28 +1962,24 @@ def test_install_os_add_params(self):
             )
 
     def test_file_copy_without_args(self):
-        ret = dict()
-        ret["message"] = "Please provide the absolute path of the file to be copied."
-        ret["out"] = False
-        self.assertEqual(junos.file_copy(), ret)
-
-    def test_file_copy_invalid_src(self):
+        self.assertRaises(TypeError, junos.file_copy)
+
+    @patch("paramiko.SSHClient")
+    @patch("scp.SCPClient.put")
+    @patch("scp.SCPClient.__init__")
+    def test_file_copy_invalid_src(self, mock_scpclient, mock_put, mock_ssh):
+        mock_scpclient.return_value = None
+        invalid_path = "invalid/file/path"
+        mock_put.side_effect = Exception(invalid_path)
         with patch("os.path.isfile") as mock_isfile:
             mock_isfile.return_value = False
             ret = dict()
-            ret["message"] = "Invalid source file path"
+            ret["message"] = 'Could not copy file : "invalid/file/path"'
             ret["out"] = False
-            self.assertEqual(junos.file_copy("invalid/file/path", "file"), ret)
+            self.assertEqual(junos.file_copy(invalid_path, "file"), ret)
 
     def test_file_copy_without_dest(self):
-        ret = dict()
-        ret[
-            "message"
-        ] = "Please provide the absolute path of the destination where the file is to be copied."
-        ret["out"] = False
-        with patch("salt.modules.junos.os.path.isfile") as mck:
-            mck.return_value = True
-            self.assertEqual(junos.file_copy("/home/user/config.set"), ret)
+        self.assertRaises(TypeError, junos.file_copy, src="/home/user/config.set")
 
     def test_file_copy(self):
         with patch("salt.modules.junos.SCP") as mock_scp, patch(
@@ -2106,16 +2182,28 @@ def test_load_none_path(self):
         self.assertEqual(ret, ret_exp)
 
     def test_load_wrong_tmp_file(self):
-        ret_exp = {"out": False, "message": "Invalid file path."}
-        with patch("salt.utils.files.mkstemp") as mock_mkstemp:
-            mock_mkstemp.return_value = "/pat/to/tmp/file"
-            ret = junos.load("/path/to/file")
-            self.assertEqual(ret, ret_exp)
+        ret_exp = {
+            "out": False,
+            "message": "Could not load configuration due to : \"[Errno 2] No such file or directory: '/pat/to/tmp/file'\"",
+            "format": "text",
+        }
+        with patch.dict(
+            junos.__salt__,
+            {
+                "cp.is_cached": MagicMock(return_value="/pat/to/tmp/file"),
+                "cp.hash_file": MagicMock(
+                    return_value={"hash_type": "sha256", "hsum": "a386e49c17"}
+                ),
+                "file.get_hash": MagicMock(return_value="a386e49c17"),
+            },
+        ):
+            with patch("os.path.getsize") as mock_getsize:
+                mock_getsize.return_value = 1000
+                ret = junos.load("salt://path/to/file")
+                self.assertEqual(ret, ret_exp)
 
     def test_load_invalid_path(self):
-        ret_exp = {"out": False, "message": "Template failed to render"}
-        ret = junos.load("/path/to/file")
-        self.assertEqual(ret, ret_exp)
+        self.assertRaises(FileNotFoundError, junos.load, path="/path/to/file")
 
     def test_load_no_extension(self):
         ret_exp = {"out": True, "message": "Successfully loaded the configuration."}
@@ -2132,17 +2220,32 @@ def test_load_no_extension(self):
             self.assertEqual(ret, ret_exp)
 
     def test_load_xml_extension(self):
+        ret_exp = {"out": True, "message": "Successfully loaded the configuration."}
+        with patch("os.path.getsize") as mock_getsize, patch(
+            "jnpr.junos.utils.config.Config.load"
+        ) as mock_load, patch("os.path.isfile") as mock_isfile:
+            mock_getsize.return_value = 1000
+            mock_isfile.return_value = True
+            ret = junos.load("/path/to/file.xml")
+            mock_load.assert_called_with(format="xml", path="/path/to/file.xml")
+            self.assertEqual(ret, ret_exp)
+
+    def test_load_xml_extension_with_kwargs(self):
         ret_exp = {"out": True, "message": "Successfully loaded the configuration."}
         with patch("os.path.getsize") as mock_getsize, patch(
             "jnpr.junos.utils.config.Config.load"
         ) as mock_load, patch("salt.utils.files.mkstemp") as mock_mkstmp, patch(
             "os.path.isfile"
-        ) as mock_isfile:
+        ) as mock_isfile, patch(
+            "salt.utils.files.fopen"
+        ) as fopen:
             mock_getsize.return_value = 1000
             mock_mkstmp.return_value = "/path/to/file"
             mock_isfile.return_value = True
-            ret = junos.load("/path/to/file.xml")
-            mock_load.assert_called_with(format="xml", path="/path/to/file")
+            ret = junos.load("/path/to/file.xml", template_vars=dict(hostname="test"))
+            mock_load.assert_called_with(
+                format="xml", path="/path/to/file", template_vars={"hostname": "test"}
+            )
             self.assertEqual(ret, ret_exp)
 
     def test_load_set_extension(self):
@@ -2156,7 +2259,7 @@ def test_load_set_extension(self):
             mock_mkstmp.return_value = "/path/to/file"
             mock_isfile.return_value = True
             ret = junos.load("/path/to/file.set")
-            mock_load.assert_called_with(format="set", path="/path/to/file")
+            mock_load.assert_called_with(format="set", path="/path/to/file.set")
             self.assertEqual(ret, ret_exp)
 
     def test_load_replace_true(self):
@@ -2263,10 +2366,17 @@ def test_get_table_wrong_path(self):
             "tablename": "ModuleTable",
             "message": "Given table file {} cannot be located".format(file),
         }
-        with patch("jnpr.junos.factory.FactoryLoader.load") as mock_load:
-            ret = junos.get_table(table, file, path)
-            self.assertEqual(ret, ret_exp)
-            mock_load.assert_not_called()
+        with patch.dict(
+            junos.__salt__, {"file.file_exists": MagicMock(return_value=False)}
+        ):
+            with patch("jnpr.junos.factory.FactoryLoader.load") as mock_load, patch(
+                "salt.utils.files.fopen"
+            ) as mock_fopen, patch(
+                "jnpr.junos.factory.FactoryLoader.load"
+            ) as mock_load:
+                ret = junos.get_table(table, file, path)
+                self.assertEqual(ret, ret_exp)
+                mock_load.assert_not_called()
 
     def test_get_table_no_path_no_file(self):
         table = "ModuleTable"
@@ -2277,13 +2387,16 @@ def test_get_table_no_path_no_file(self):
             "tablename": "ModuleTable",
             "message": "Given table file {} cannot be located".format(file),
         }
-        with patch("jnpr.junos.factory.FactoryLoader.load") as mock_load, patch(
-            "glob.glob"
-        ) as mock_fopen:
-            mock_fopen.return_value = []
-            ret = junos.get_table(table, file)
-            self.assertEqual(ret, ret_exp)
-            mock_load.assert_not_called()
+        with patch.dict(
+            junos.__salt__, {"file.file_exists": MagicMock(return_value=False)}
+        ):
+            with patch("jnpr.junos.factory.FactoryLoader.load") as mock_load, patch(
+                "glob.glob"
+            ) as mock_fopen:
+                mock_fopen.return_value = []
+                ret = junos.get_table(table, file)
+                self.assertEqual(ret, ret_exp)
+                mock_load.assert_not_called()
 
     def test_get_table_yaml_load_error(self):
         table = "ModuleTable"

From 24498a817a6c9f81e05c27d444dd576f5bb9e123 Mon Sep 17 00:00:00 2001
From: Nitin Kumar <nitinkr@juniper.net>
Date: Wed, 20 May 2020 22:49:49 +0530
Subject: [PATCH 3/6] fix isort

---
 salt/modules/junos.py            | 2 +-
 tests/unit/modules/test_junos.py | 3 +--
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/salt/modules/junos.py b/salt/modules/junos.py
index f3622e5050e..55107200e0d 100644
--- a/salt/modules/junos.py
+++ b/salt/modules/junos.py
@@ -21,6 +21,7 @@
 import json
 import logging
 import os
+import tempfile
 import traceback
 from functools import wraps
 
@@ -31,7 +32,6 @@
 import salt.utils.stringutils
 import yaml
 from salt.ext import six
-import tempfile
 
 try:
     from lxml import etree
diff --git a/tests/unit/modules/test_junos.py b/tests/unit/modules/test_junos.py
index 1a937d8e208..62f282acd6b 100644
--- a/tests/unit/modules/test_junos.py
+++ b/tests/unit/modules/test_junos.py
@@ -10,10 +10,9 @@
 # Import salt modules
 import salt.modules.junos as junos
 from salt.ext import six
-
 # Import test libs
 from tests.support.mixins import LoaderModuleMockMixin, XMLEqualityMixin
-from tests.support.mock import ANY, PropertyMock, call, mock_open, patch, MagicMock
+from tests.support.mock import ANY, MagicMock, PropertyMock, call, mock_open, patch
 from tests.support.unit import TestCase, skipIf
 
 # Import 3rd-party libs

From 52b577b74c2bd942306a7086e11d289b54923e48 Mon Sep 17 00:00:00 2001
From: Nitin Kumar <nitinkr@juniper.net>
Date: Wed, 20 May 2020 22:56:01 +0530
Subject: [PATCH 4/6] fix isort

---
 tests/unit/modules/test_junos.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/unit/modules/test_junos.py b/tests/unit/modules/test_junos.py
index 62f282acd6b..799e40830cc 100644
--- a/tests/unit/modules/test_junos.py
+++ b/tests/unit/modules/test_junos.py
@@ -10,6 +10,7 @@
 # Import salt modules
 import salt.modules.junos as junos
 from salt.ext import six
+
 # Import test libs
 from tests.support.mixins import LoaderModuleMockMixin, XMLEqualityMixin
 from tests.support.mock import ANY, MagicMock, PropertyMock, call, mock_open, patch

From a6b44f0a902dc865c9bd8a4be4895ddb88b54795 Mon Sep 17 00:00:00 2001
From: Nitin Kumar <nitinkr@juniper.net>
Date: Wed, 20 May 2020 23:46:21 +0530
Subject: [PATCH 5/6] added comment for code flow

---
 salt/modules/junos.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/salt/modules/junos.py b/salt/modules/junos.py
index 55107200e0d..f21d814888e 100644
--- a/salt/modules/junos.py
+++ b/salt/modules/junos.py
@@ -94,8 +94,10 @@ def __enter__(self):
                 proxy_hash = __salt__["file.get_hash"](local_cache_path)
                 # check if hash is same, else copy newly
                 if master_hash.get("hsum") == proxy_hash:
+                    # kwargs will have values when path is a template
                     if self._kwargs:
                         self._cached_file = salt.utils.files.mkstemp()
+                        # local copy is a template, hence need to render
                         with salt.utils.files.fopen(self._cached_file, "w") as fp:
                             template_string = __salt__["slsutil.renderer"](
                                 path=local_cache_path,

From f308a577738f4e70d27927594cf3c26bab79b4b8 Mon Sep 17 00:00:00 2001
From: Nitin Kumar <nitinkr@juniper.net>
Date: Thu, 21 May 2020 00:22:52 +0530
Subject: [PATCH 6/6] keep __virtual__ function on top

---
 salt/modules/junos.py | 34 +++++++++++++++++-----------------
 1 file changed, 17 insertions(+), 17 deletions(-)

diff --git a/salt/modules/junos.py b/salt/modules/junos.py
index f21d814888e..e38ee0cf1ff 100644
--- a/salt/modules/junos.py
+++ b/salt/modules/junos.py
@@ -71,6 +71,23 @@
 __proxyenabled__ = ["junos"]
 
 
+def __virtual__():
+    """
+    We need the Junos adapter libraries for this
+    module to work.  We also need a proxymodule entry in __opts__
+    in the opts dictionary
+    """
+    if HAS_JUNOS and "proxy" in __opts__:
+        return __virtualname__
+    else:
+        return (
+            False,
+            "The junos or dependent module could not be loaded: "
+            "junos-eznc or jxmlease or yamlordereddictloader or "
+            "proxy could not be loaded.",
+        )
+
+
 class HandleFileCopy:
     """
     To figure out proper path either from proxy local file system
@@ -148,23 +165,6 @@ def __exit__(self, exc_type, exc_value, exc_traceback):
             log.debug("Deleted cached folder: {0}".format(self._cached_folder))
 
 
-def __virtual__():
-    """
-    We need the Junos adapter libraries for this
-    module to work.  We also need a proxymodule entry in __opts__
-    in the opts dictionary
-    """
-    if HAS_JUNOS and "proxy" in __opts__:
-        return __virtualname__
-    else:
-        return (
-            False,
-            "The junos or dependent module could not be loaded: "
-            "junos-eznc or jxmlease or yamlordereddictloader or "
-            "proxy could not be loaded.",
-        )
-
-
 def timeoutDecorator(function):
     @wraps(function)
     def wrapper(*args, **kwargs):