Add JunOS pre-transfer free-space check (NAPPS-1085)

changes/373.added

@@ -0,0 +1 @@
+Added a pre-transfer free-space check to Juniper JunOS ``file_copy`` and ``remote_file_copy`` that raises ``NotEnoughFreeSpaceError`` when the target filesystem lacks room for the image.

pyntc/devices/jnpr_device.py

@@ -6,6 +6,7 @@
 import time
 import warnings
 from tempfile import NamedTemporaryFile
+from urllib.parse import urlparse
 
 from jnpr.junos import Device as JunosNativeDevice
 from jnpr.junos.exception import ConfigLoadError
@@ -18,9 +19,54 @@
 from pyntc import log
 from pyntc.devices.base_device import BaseDevice, fix_docs
 from pyntc.devices.tables.jnpr.loopback import LoopbackTable  # pylint: disable=no-name-in-module
-from pyntc.errors import CommandError, CommandListError, FileTransferError, OSInstallError, RebootTimeoutError
+from pyntc.errors import (
+    CommandError,
+    CommandListError,
+    FileSystemNotFoundError,
+    FileTransferError,
+    OSInstallError,
+    RebootTimeoutError,
+)
 from pyntc.utils.models import FileCopyModel
 
+# Multipliers for Junos ``df``-style size suffixes. Junos formats available
+# space with binary (1024-based) units in its ``<available-blocks format="...">``
+# XML attribute (e.g., "126M", "1.0G").
+_JUNOS_SIZE_UNIT_MULTIPLIERS = {
+    "": 1,
+    "B": 1,
+    "K": 1024,
+    "M": 1024**2,
+    "G": 1024**3,
+    "T": 1024**4,
+    "P": 1024**5,
+}
+_JUNOS_AVAIL_FORMAT_RE = re.compile(r"^\s*(\d+(?:\.\d+)?)\s*([BKMGTP]?)\s*$", re.IGNORECASE)
+# Default mount point to probe when callers do not specify one. ``/var/tmp`` is
+# the standard destination for ``fs.cp`` transfers on Junos (remote device
+# mount point, not a local temp directory).
+_JUNOS_DEFAULT_FILE_SYSTEM = "/var/tmp"  # noqa: S108
+
+# Hashing algorithms that Junos implements for the ``file checksum`` RPC.
+# Junos does NOT implement sha512; callers passing it will be rejected at the
+# driver boundary rather than surfacing PyEZ's raw ValueError deeper in the
+# stack. Mirrors the pattern used by EOS and NXOS drivers.
+JUNOS_SUPPORTED_HASHING_ALGORITHMS = {"md5", "sha1", "sha256"}
+
+
+def _mount_encloses_path(mount, path):
+    """Return True if ``mount`` is the filesystem that contains ``path``.
+
+    Matches with directory-boundary semantics (the same rule ``df`` uses) so
+    ``/vari`` is not mistaken for a prefix of ``/var/tmp``. ``/`` encloses
+    every path.
+    """
+    if mount == "/":
+        return True
+    if path == mount:
+        return True
+    return path.startswith(mount.rstrip("/") + "/")
+
 
 @fix_docs
 class JunosDevice(BaseDevice):
@@ -62,6 +108,85 @@ def _file_copy_local_md5(self, filepath, blocksize=2**20):
                     buf = file_name.read(blocksize)
             return md5_hash.hexdigest()
 
+    def _get_free_space(self, file_system=None):
+        """Return free bytes on the filesystem containing ``file_system``.
+
+        Probes the device via ``get-system-storage-information`` (invoked by
+        PyEZ ``FS.storage_usage``) and parses the human-readable
+        ``available-blocks`` ``format`` attribute (e.g., ``"126M"``, ``"1.0G"``)
+        into bytes. The human-readable string is used rather than the raw
+        block count because PyEZ does not expose a native block size and
+        Junos block semantics can vary by release.
+
+        ``file_system`` is resolved by **longest-prefix mount match** — the
+        same logic ``df`` uses — so a caller asking about ``/var/tmp`` on a
+        platform that only mounts ``/var`` (e.g., SRX hardware) still gets
+        back the correct filesystem's free space. ``/`` is always a fallback
+        when nothing more specific matches.
+
+        Args:
+            file_system (str, optional): Target path. When ``None`` (the
+                default), the probe uses ``_JUNOS_DEFAULT_FILE_SYSTEM``
+                (``/var/tmp`` — the standard destination for ``fs.cp`` copies
+                on Junos).
+
+        Returns:
+            int: Free bytes available on the resolved filesystem.
+
+        Raises:
+            FileSystemNotFoundError: When no mount point encloses ``file_system``
+                (i.e., not even ``/`` is present in ``storage_usage``).
+            CommandError: When the ``avail`` format string cannot be parsed.
+        """
+        if file_system is None:
+            file_system = _JUNOS_DEFAULT_FILE_SYSTEM
+
+        usage = self.fs.storage_usage()
+        best_info = None
+        best_mount = None
+        best_len = -1
+        for _dev, info in usage.items():
+            mount = info.get("mount")
+            if not mount or not _mount_encloses_path(mount, file_system):
+                continue
+            if len(mount) > best_len:
+                best_info = info
+                best_mount = mount
+                best_len = len(mount)
+
+        if best_info is None:
+            log.error(
+                "Host %s: no mount encloses %s in storage_usage output.",
+                self.host,
+                file_system,
+            )
+            raise FileSystemNotFoundError(hostname=self.host, command="show system storage")
+
+        avail = best_info.get("avail", "")
+        match = _JUNOS_AVAIL_FORMAT_RE.match(str(avail))
+        if match is None:
+            log.error(
+                "Host %s: could not parse avail %r for mount %s.",
+                self.host,
+                avail,
+                best_mount,
+            )
+            raise CommandError(
+                command="show system storage",
+                message=f"Unable to parse available space {avail!r} for {best_mount}.",
+            )
+        size = float(match.group(1))
+        multiplier = _JUNOS_SIZE_UNIT_MULTIPLIERS[match.group(2).upper()]
+        free_bytes = int(size * multiplier)
+        log.debug(
+            "Host %s: %s bytes free on %s (resolved from %s).",
+            self.host,
+            free_bytes,
+            best_mount,
+            file_system,
+        )
+        return free_bytes
+
     def _get_interfaces(self):
         eth_ifaces = EthPortTable(self.native)
         eth_ifaces.get()
@@ -304,11 +429,15 @@ def file_copy(self, src, dest=None, **kwargs):
 
         Raises:
             FileTransferError: Raised when unable to verify file was transferred succesfully.
+            NotEnoughFreeSpaceError: When the target filesystem has fewer free bytes
+                than ``src`` requires.
         """
         if not self.file_copy_remote_exists(src, dest, **kwargs):
             if dest is None:
                 dest = os.path.basename(src)
 
+            self._check_free_space(os.path.getsize(src))
+
             with SCP(self.native) as scp:
                 scp.put(src, remote_path=dest)
 
@@ -505,11 +634,21 @@ def get_remote_checksum(self, filename, hashing_algorithm="md5"):
 
         Args:
             filename (str): The name of the file to check for on the remote device.
-            hashing_algorithm (str): The hashing algorithm to use. Valid values are 'md5', 'sha1', and 'sha256'. Defaults to 'md5'.
+            hashing_algorithm (str): The hashing algorithm to use. Valid values are
+                those in ``JUNOS_SUPPORTED_HASHING_ALGORITHMS`` (``md5``, ``sha1``,
+                ``sha256``). Defaults to ``md5``.
 
         Returns:
             (str): The checksum of the remote file or None if the file is not found.
+
+        Raises:
+            ValueError: When ``hashing_algorithm`` is not one Junos implements.
         """
+        if hashing_algorithm.lower() not in JUNOS_SUPPORTED_HASHING_ALGORITHMS:
+            raise ValueError(
+                f"Unsupported hashing algorithm '{hashing_algorithm}' for Junos. "
+                f"Supported algorithms: {sorted(JUNOS_SUPPORTED_HASHING_ALGORITHMS)}"
+            )
         return self.fs.checksum(path=filename, calc=hashing_algorithm)
 
     def compare_file_checksum(self, checksum, filename, hashing_algorithm="md5"):
@@ -525,24 +664,41 @@ def compare_file_checksum(self, checksum, filename, hashing_algorithm="md5"):
         """
         return checksum == self.get_remote_checksum(filename, hashing_algorithm)
 
-    def remote_file_copy(self, src: FileCopyModel = None, dest=None):
+    def remote_file_copy(self, src: FileCopyModel = None, dest=None, file_system: str | None = None, **kwargs):
         """Copy a file to a remote device.
 
         Args:
             src (FileCopyModel): The source file model.
             dest (str): The destination file path on the remote device.
+            file_system (str, optional): Mount point used for the pre-transfer
+                free-space check. When ``None`` (the default), the probe uses
+                ``_JUNOS_DEFAULT_FILE_SYSTEM`` (``/var/tmp``).
+            **kwargs (Any): Accepted for parity with ``BaseDevice.remote_file_copy``;
+                other drivers may forward extra options.
 
         Raises:
             TypeError: If src is not an instance of FileCopyModel.
             FileTransferError: If there is an error during file transfer or if the file cannot be verified after transfer.
+            NotEnoughFreeSpaceError: If ``src.file_size_bytes`` is set and the
+                target mount point has fewer free bytes than ``src.file_size_bytes``.
+                When ``file_size`` is omitted from ``src`` the pre-transfer space
+                check is skipped entirely.
         """
         if not isinstance(src, FileCopyModel):
             raise TypeError("src must be an instance of FileCopyModel")
 
         if self.verify_file(src.checksum, dest, hashing_algorithm=src.hashing_algorithm):
             return
 
-        if not self.fs.cp(from_path=src.download_url, to_path=dest, dev_timeout=src.timeout):
+        self._pre_transfer_space_check(src, file_system=file_system)
+
+        # Junos ``fs.cp`` requires the filename in the URL; append ``src.file_name``
+        # when the URL carries no path so callers can point at a bare host.
+        source_url = src.download_url
+        if not urlparse(source_url).path.strip("/"):
+            source_url = f"{source_url.rstrip('/')}/{src.file_name}"
+
+        if not self.fs.cp(from_path=source_url, to_path=dest, dev_timeout=src.timeout):
             raise FileTransferError(message=f"Unable to copy file from remote url {src.clean_url}")
 
         # Some devices take a while to sync the filesystem after a copy but netconf returns before the sync completes

tests/integration/_helpers.py

@@ -22,11 +22,15 @@
 def build_file_copy_model(url_env_var):
     """Build a ``FileCopyModel`` from a per-protocol URL env var.
 
-    Calls ``pytest.skip`` if the URL, ``FILE_CHECKSUM``, or ``FILE_SIZE`` env
-    vars are not set.
+    Reads ``FILE_HASH_ALGO`` and ``FILE_CHECKSUM`` from the environment.
+    An autouse fixture in ``conftest.py`` sets both to the running test
+    module's platform default before each module runs, so individual
+    tests never have to hardcode an algorithm. Calls ``pytest.skip`` when
+    any required env var is missing.
     """
     url = os.environ.get(url_env_var)
     checksum = os.environ.get("FILE_CHECKSUM")
+    hashing_algorithm = os.environ.get("FILE_HASH_ALGO", "sha512")
     file_name = os.environ.get("FILE_NAME") or (posixpath.basename(url.split("?")[0]) if url else None)
     file_size = int(os.environ.get("FILE_SIZE", "0"))
     file_size_unit = os.environ.get("FILE_SIZE_UNIT", "bytes")
@@ -40,7 +44,7 @@ def build_file_copy_model(url_env_var):
         file_name=file_name,
         file_size=file_size,
         file_size_unit=file_size_unit,
-        hashing_algorithm="sha512",
+        hashing_algorithm=hashing_algorithm,
         timeout=900,
     )
 

tests/integration/conftest.py

@@ -9,6 +9,76 @@
 
 from ._helpers import PROTOCOL_URL_VARS
 
+# Each driver's integration test module is mapped to the hashing algorithm
+# its device family implements. ``conftest.py`` owns this mapping (rather
+# than per-module constants or .env entries) so the user can run every
+# driver's integration suite in one ``pytest tests/integration`` invocation
+# and each module automatically picks up the algorithm its hardware
+# supports. Junos does not implement sha512 — sha256 is typical on SRX /
+# MX. EOS and ASA both implement sha512. Extend the map as new drivers get
+# integration tests.
+_PLATFORM_HASH_ALGOS = {
+    "test_eos_device": "sha512",
+    "test_asa_device": "sha512",
+    "test_jnpr_device": "sha256",
+    "test_ios_device": "md5",
+    "test_nxos_device": "sha256",
+}
+
+# Maps each hashing algorithm to the suffix convention used on the
+# per-algorithm checksum env vars (``FILE_CHECKSUM_512`` / ``_256`` /
+# ``_MD5``). The user exports one checksum per algorithm once and
+# ``_configure_integration_env`` copies the right one into
+# ``FILE_CHECKSUM`` before the module runs.
+_HASH_ALGO_ENV_SUFFIXES = {"sha512": "512", "sha256": "256", "md5": "MD5"}
+
+
+@pytest.fixture(scope="module", autouse=True)
+def _configure_integration_env(request):
+    """Set ``FILE_HASH_ALGO`` / ``FILE_CHECKSUM`` per test module.
+
+    Resolves the module-specific hashing algorithm from
+    ``_PLATFORM_HASH_ALGOS``, copies the matching ``FILE_CHECKSUM_<SUFFIX>``
+    env var into ``FILE_CHECKSUM``, and restores any prior values when
+    the module finishes. Test files not listed in the map are left alone —
+    they either carry no hashing dependency or set their own env.
+    """
+    module_name = request.module.__name__.split(".")[-1]
+    algo = _PLATFORM_HASH_ALGOS.get(module_name)
+    if algo is None:
+        yield
+        return
+
+    suffix = _HASH_ALGO_ENV_SUFFIXES.get(algo)
+    checksum = os.environ.get(f"FILE_CHECKSUM_{suffix}") if suffix else None
+
+    # Skip the env overwrite entirely when the suffix env var is missing.
+    # Otherwise we would pin FILE_HASH_ALGO to the platform's algo while
+    # leaving FILE_CHECKSUM inherited from the shell (possibly a different
+    # algo's hash) — tests would then fail at verify with a confusing
+    # mismatch. Leaving both vars alone lets the user's own pair win, or
+    # lets ``build_file_copy_model`` skip cleanly on missing env.
+    if checksum is None:
+        yield
+        return
+
+    prior_algo = os.environ.get("FILE_HASH_ALGO")
+    prior_checksum = os.environ.get("FILE_CHECKSUM")
+
+    os.environ["FILE_HASH_ALGO"] = algo
+    os.environ["FILE_CHECKSUM"] = checksum
+
+    yield
+
+    if prior_algo is None:
+        os.environ.pop("FILE_HASH_ALGO", None)
+    else:
+        os.environ["FILE_HASH_ALGO"] = prior_algo
+    if prior_checksum is None:
+        os.environ.pop("FILE_CHECKSUM", None)
+    else:
+        os.environ["FILE_CHECKSUM"] = prior_checksum
+
 
 @pytest.fixture(scope="module")
 def any_file_copy_model():
@@ -19,6 +89,7 @@ def any_file_copy_model():
     protocol URL / ``FILE_CHECKSUM`` / ``FILE_SIZE`` env vars are set.
     """
     checksum = os.environ.get("FILE_CHECKSUM")
+    checksum_algo = os.environ.get("FILE_HASH_ALGO", "sha512")
     file_size = int(os.environ.get("FILE_SIZE", "0"))
     file_size_unit = os.environ.get("FILE_SIZE_UNIT", "bytes")
     for env_var in PROTOCOL_URL_VARS.values():
@@ -31,7 +102,7 @@ def any_file_copy_model():
                 file_name=file_name,
                 file_size=file_size,
                 file_size_unit=file_size_unit,
-                hashing_algorithm="sha512",
+                hashing_algorithm=checksum_algo,
                 timeout=900,
             )
     pytest.skip("No protocol URL / FILE_CHECKSUM / FILE_SIZE environment variables not set")