Add EOS pre-transfer free-space check (NAPPS-1091)

changes/370.added

@@ -0,0 +1,2 @@
+Added a pre-transfer free-space check to EOS ``file_copy`` and ``remote_file_copy`` that raises ``NotEnoughFreeSpaceError`` when the target filesystem lacks room for the image.
+Added ``file_size_unit`` (``bytes``, ``megabytes``, or ``gigabytes``; default ``bytes``) and a computed ``file_size_bytes`` to ``FileCopyModel`` so ``remote_file_copy`` can verify free space against a caller-supplied size; when ``file_size`` is omitted the pre-transfer check is skipped.

docs/user/lib_getting_started.md

@@ -264,7 +264,12 @@ from pyntc.utils.models import FileCopyModel
 ...     checksum='abc123def456',
 ...     hashing_algorithm='md5',
 ...     file_name='newconfig.cfg',
-        vrf='Mgmt-vrf'
+...     # file_size is optional. When supplied, remote_file_copy verifies
+...     # the target device has room before starting the transfer. When
+...     # omitted, the pre-transfer space check is skipped.
+...     file_size=512,
+...     file_size_unit='megabytes',
+...     vrf='Mgmt-vrf',
 ... )
 >>> for device in devices:
 ...    device.remote_file_copy(source_file)

pyntc/devices/base_device.py

@@ -2,11 +2,14 @@
 
 import hashlib
 import importlib
+import logging
 import warnings
 
-from pyntc.errors import FeatureNotFoundError, NTCError
+from pyntc.errors import FeatureNotFoundError, NotEnoughFreeSpaceError, NTCError
 from pyntc.utils.models import FileCopyModel
 
+log = logging.getLogger(__name__)
+
 
 def fix_docs(cls):
     """Create docstring at runtime.
@@ -52,6 +55,90 @@ def __init__(self, host, username, password, device_type=None, **kwargs):  # noq
         self._model = None
         self._vlans = None
 
+    def _check_free_space(self, required_bytes, file_system=None):
+        """Raise NotEnoughFreeSpaceError when the target filesystem lacks room for a transfer.
+
+        Drivers call this from ``file_copy`` and ``remote_file_copy`` before starting a
+        transfer. The concrete per-platform probe is ``_get_free_space``; this helper
+        centralises the comparison, logging, and error shape so every driver behaves the
+        same way.
+
+        Args:
+            required_bytes (int): Number of bytes the pending transfer needs.
+            file_system (str, optional): Target filesystem passed through to
+                ``_get_free_space``. Drivers that auto-detect a filesystem should resolve
+                it before calling.
+
+        Raises:
+            NotEnoughFreeSpaceError: When the device reports fewer free bytes than
+                ``required_bytes``.
+        """
+        available = self._get_free_space(file_system)
+        log.debug(
+            "Host %s: filesystem %s has %s bytes free; %s bytes required.",
+            self.host,
+            file_system,
+            available,
+            required_bytes,
+        )
+        if available < required_bytes:
+            log.error(
+                "Host %s: insufficient free space on %s (%s free, %s required).",
+                self.host,
+                file_system,
+                available,
+                required_bytes,
+            )
+            raise NotEnoughFreeSpaceError(
+                hostname=self.host,
+                required=required_bytes,
+                available=available,
+                file_system=file_system,
+            )
+
+    def _get_free_space(self, file_system=None):
+        """Return the number of free bytes on ``file_system``.
+
+        Drivers override this to issue a platform-specific command (e.g., ``dir`` on
+        EOS/IOS/NXOS, ``show system storage`` on JUNOS) and parse the result. Drivers
+        whose platform exposes only one practical filesystem (e.g., EOS) may safely
+        ignore ``file_system``.
+
+        Args:
+            file_system (str, optional): The target filesystem to inspect. Drivers
+                that support multiple filesystems should treat ``None`` as "auto-detect
+                the default filesystem".
+
+        Returns:
+            int: Free bytes available on ``file_system``.
+
+        Raises:
+            NotImplementedError: When a driver has not yet implemented the probe.
+        """
+        raise NotImplementedError(f"{self.__class__.__name__} does not implement _get_free_space")
+
+    def _pre_transfer_space_check(self, src, file_system=None):
+        """Run the free-space check if ``src.file_size_bytes`` is populated.
+
+        Drivers call this from ``remote_file_copy`` so the check is skipped
+        (fail-open) when the caller omits ``FileCopyModel.file_size``; when set,
+        ``_check_free_space`` raises ``NotEnoughFreeSpaceError`` if the device
+        lacks room. Lives on ``BaseDevice`` so every driver inherits the same
+        shape without duplicating the if/else.
+
+        Args:
+            src (FileCopyModel): The source specification for the pending transfer.
+            file_system (str, optional): Target filesystem; passed through to
+                ``_check_free_space`` (and from there to ``_get_free_space``).
+        """
+        if src.file_size_bytes is not None:
+            self._check_free_space(src.file_size_bytes, file_system=file_system)
+        else:
+            log.debug(
+                "Host %s: no file_size on FileCopyModel; skipping pre-transfer space check.",
+                self.host,
+            )
+
     def _image_booted(self, image_name, **vendor_specifics):
         """Determine if a particular image is serving as the active OS.
 

pyntc/devices/eos_device.py

@@ -88,6 +88,34 @@ def _file_copy_instance(self, src, dest=None, file_system="/mnt/flash"):
         log.debug("Host %s: File copy instance %s.", self.host, file_copy)
         return file_copy
 
+    def _get_free_space(self, file_system=None):  # pylint: disable=unused-argument
+        """Return free bytes on ``file_system`` as reported by Arista's ``dir`` output.
+
+        EOS only exposes a single flash filesystem in practice, and ``dir`` always
+        prints ``<total> bytes total (<free> bytes free)`` as the last line, so the
+        ``file_system`` argument is accepted for API parity with other drivers but is
+        otherwise unused.
+
+        Args:
+            file_system (str, optional): Ignored; retained for BaseDevice API parity.
+
+        Returns:
+            int: Free bytes available on the target filesystem.
+
+        Raises:
+            CommandError: When the ``dir`` output does not contain a parseable
+                ``(N bytes free)`` trailer.
+        """
+        raw_data = self.show("dir", raw_text=True)
+        # Example: 3634421760 bytes total (1951289344 bytes free)
+        match = re.search(r"\((\d+)\s+bytes\s+free\)", raw_data)
+        if match is None:
+            log.error("Host %s: could not parse free space from 'dir' output.", self.host)
+            raise CommandError(command="dir", message="Unable to parse free space from dir output.")
+        free_bytes = int(match.group(1))
+        log.debug("Host %s: %s bytes free on flash.", self.host, free_bytes)
+        return free_bytes
+
     def _get_file_system(self):
         """Determine the default file system or directory for device.
 
@@ -371,6 +399,8 @@ def file_copy(self, src, dest=None, file_system=None):
 
         Raises:
             FileTransferError: raise exception if there is an error
+            NotEnoughFreeSpaceError: When ``file_system`` has fewer free bytes than
+                ``src`` requires.
         """
         self.open()
         self.enable()
@@ -379,6 +409,7 @@ def file_copy(self, src, dest=None, file_system=None):
             file_system = self._get_file_system()
 
         if not self.file_copy_remote_exists(src, dest, file_system):
+            self._check_free_space(os.path.getsize(src), file_system=file_system)
             file_copy = self._file_copy_instance(src, dest, file_system=file_system)
 
             try:
@@ -584,6 +615,10 @@ def remote_file_copy(self, src: FileCopyModel, dest: str | None = None, file_sys
             ValueError: If the URL scheme is unsupported or URL contains query strings.
             FileTransferError: If transfer or verification fails.
             FileSystemNotFoundError: If filesystem cannot be determined.
+            NotEnoughFreeSpaceError: If ``src.file_size_bytes`` is set and
+                ``file_system`` 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")
@@ -606,6 +641,8 @@ def remote_file_copy(self, src: FileCopyModel, dest: str | None = None, file_sys
         self.open()
         self.enable()
 
+        self._pre_transfer_space_check(src, file_system)
+
         if src.scheme == "tftp" or src.username is None:
             command, detect_prompt = self._build_url_copy_command_simple(src, file_system, dest)
         else:

pyntc/errors.py

@@ -206,15 +206,23 @@ def __init__(self, hostname, wait_time):
 class NotEnoughFreeSpaceError(NTCError):
     """Error for not having enough free space to transfer a file."""
 
-    def __init__(self, hostname, min_space):
+    def __init__(self, hostname, min_space=None, *, required=None, available=None, file_system=None):
         """
         Error for not having enough free space to transfer a file.
 
         Args:
-            hostname (str): The hostname of the device that did not boot back up.
-            min_space (str): The minimum amount of space required to transfer the file.
-        """
-        message = f"{hostname} does not meet the minimum disk space requirements of {min_space}"
+            hostname (str): The hostname of the device being checked.
+            min_space (str, optional): The minimum amount of space required. Retained for
+                backward compatibility with callers that only know the required value.
+            required (int, optional): Required bytes for the pending transfer.
+            available (int, optional): Free bytes currently available on the target filesystem.
+            file_system (str, optional): The target filesystem that was checked.
+        """
+        if required is not None and available is not None:
+            location = f"{file_system} " if file_system else ""
+            message = f"{hostname}: {location}has {available} bytes free; {required} bytes required for transfer"
+        else:
+            message = f"{hostname} does not meet the minimum disk space requirements of {min_space}"
         super().__init__(message)
 
 

pyntc/utils/models.py

@@ -7,6 +7,14 @@
 # Use Hashing algorithms from Nautobot's supported list.
 HASHING_ALGORITHMS = {"md5", "sha1", "sha224", "sha384", "sha256", "sha512", "sha3", "blake2", "blake3"}
 
+# Supported units for FileCopyModel.file_size, mapped to their multiplier in bytes.
+# Conversions use binary units (1 MB = 1024**2 bytes) to match network-device reporting.
+FILE_SIZE_UNITS = {
+    "bytes": 1,
+    "megabytes": 1024**2,
+    "gigabytes": 1024**3,
+}
+
 
 @dataclass
 class FileCopyModel:
@@ -16,9 +24,10 @@ class FileCopyModel:
         download_url (str): The URL to download the file from. Can include credentials, but it's recommended to use the username and token fields instead for security reasons.
         checksum (str): The expected checksum of the file.
         file_name (str): The name of the file to be saved on the device.
+        file_size (int, optional): The expected size of the file. When supplied, ``remote_file_copy`` verifies the target device has room before starting the transfer. When omitted, the pre-transfer space check is skipped (callers can probe the source URL themselves and populate this field). Defaults to ``None``.
+        file_size_unit (str, optional): Unit that ``file_size`` is expressed in. One of ``"bytes"``, ``"megabytes"``, ``"gigabytes"``. Only consulted when ``file_size`` is supplied. Defaults to ``"bytes"``.
         hashing_algorithm (str, optional): The hashing algorithm to use for checksum verification. Defaults to "md5".
         timeout (int, optional): The timeout for the download operation in seconds. Defaults to 900.
-        file_size (int, optional): The expected size of the file in bytes. Optional but can be used for an additional layer of verification.
         username (str, optional): The username for authentication if required by the URL. Optional if credentials are included in the URL.
         token (str, optional): The password or token for authentication if required by the URL. Optional if credentials are included in the URL.
         vrf (str, optional): The VRF to use for the download if the device supports VRFs. Optional.
@@ -28,26 +37,40 @@ class FileCopyModel:
     download_url: str
     checksum: str
     file_name: str
+    file_size: Optional[int] = None
+    file_size_unit: str = "bytes"
     hashing_algorithm: str = "md5"
     timeout: int = 900  # Timeout for the download operation in seconds
-    file_size: Optional[int] = None  # Size in bytes
     username: Optional[str] = None
     token: Optional[str] = None  # Password/Token
     vrf: Optional[str] = None
     ftp_passive: bool = True
 
-    # Computed fields derived from download_url — not passed to the constructor
+    # Computed fields derived from download_url and file_size — not passed to the constructor
     clean_url: str = field(init=False)
     scheme: str = field(init=False)
     hostname: str = field(init=False)
     port: Optional[int] = field(init=False)
     path: str = field(init=False)
+    file_size_bytes: Optional[int] = field(init=False)
 
     def __post_init__(self):
         """Validate the input and prepare the clean URL after initialization."""
         if self.hashing_algorithm.lower() not in HASHING_ALGORITHMS:
             raise ValueError(f"Unsupported algorithm. Choose from: {HASHING_ALGORITHMS}")
 
+        unit = self.file_size_unit.lower()
+        if unit not in FILE_SIZE_UNITS:
+            raise ValueError(f"Unsupported file_size_unit. Choose from: {sorted(FILE_SIZE_UNITS)}")
+        self.file_size_unit = unit
+
+        if self.file_size is None:
+            self.file_size_bytes = None
+        else:
+            if self.file_size < 0:
+                raise ValueError("file_size must be a non-negative integer.")
+            self.file_size_bytes = self.file_size * FILE_SIZE_UNITS[unit]
+
         parsed = urlparse(self.download_url)
 
         # Extract username/password from URL if not already provided as arguments

tests/integration/_helpers.py

@@ -0,0 +1,57 @@
+"""Shared helpers for integration tests that drive ``remote_file_copy``."""
+
+import os
+import posixpath
+
+import pytest
+
+from pyntc.utils.models import FileCopyModel
+
+# Every protocol that ``remote_file_copy`` might transfer from. Individual device
+# test modules can narrow this set when they only support a subset.
+PROTOCOL_URL_VARS = {
+    "ftp": "FTP_URL",
+    "tftp": "TFTP_URL",
+    "scp": "SCP_URL",
+    "http": "HTTP_URL",
+    "https": "HTTPS_URL",
+    "sftp": "SFTP_URL",
+}
+
+
+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.
+    """
+    url = os.environ.get(url_env_var)
+    checksum = os.environ.get("FILE_CHECKSUM")
+    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")
+
+    if not all([url, checksum, file_name, file_size]):
+        pytest.skip(f"{url_env_var} / FILE_CHECKSUM / FILE_SIZE environment variables not set")
+
+    return FileCopyModel(
+        download_url=url,
+        checksum=checksum,
+        file_name=file_name,
+        file_size=file_size,
+        file_size_unit=file_size_unit,
+        hashing_algorithm="sha512",
+        timeout=900,
+    )
+
+
+def first_available_url(protocol_url_vars=None):
+    """Return ``(scheme, url)`` for the first configured protocol URL.
+
+    ``(None, None)`` when none of the env vars in ``protocol_url_vars`` are set.
+    """
+    for scheme, env_var in (protocol_url_vars or PROTOCOL_URL_VARS).items():
+        url = os.environ.get(env_var)
+        if url:
+            return scheme, url
+    return None, None

tests/integration/conftest.py

@@ -0,0 +1,37 @@
+"""Shared fixtures for pyntc integration tests."""
+
+import os
+import posixpath
+
+import pytest
+
+from pyntc.utils.models import FileCopyModel
+
+from ._helpers import PROTOCOL_URL_VARS
+
+
+@pytest.fixture(scope="module")
+def any_file_copy_model():
+    """Return a ``FileCopyModel`` using the first available protocol URL.
+
+    Used by tests that only need a file reference (existence checks, checksum
+    verification) without caring about the transfer protocol. Skips if no
+    protocol URL / ``FILE_CHECKSUM`` / ``FILE_SIZE`` env vars are set.
+    """
+    checksum = os.environ.get("FILE_CHECKSUM")
+    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():
+        url = os.environ.get(env_var)
+        if url and checksum and file_size:
+            file_name = os.environ.get("FILE_NAME") or posixpath.basename(url.split("?")[0])
+            return FileCopyModel(
+                download_url=url,
+                checksum=checksum,
+                file_name=file_name,
+                file_size=file_size,
+                file_size_unit=file_size_unit,
+                hashing_algorithm="sha512",
+                timeout=900,
+            )
+    pytest.skip("No protocol URL / FILE_CHECKSUM / FILE_SIZE environment variables not set")