rimport: stage_data() now checks whether file can be downloaded.

Author: samsrabin

rimport

@@ -15,13 +15,16 @@ import shutil
 import sys
 from pathlib import Path
 from typing import Iterable, List
+from urllib.request import Request, urlopen
+from urllib.error import HTTPError
 
 import shared
 
 DEFAULT_INPUTDATA_ROOT = Path(shared.DEFAULT_INPUTDATA_ROOT)
 DEFAULT_STAGING_ROOT = Path(shared.DEFAULT_STAGING_ROOT)
 STAGE_OWNER = "cesmdata"
 INDENT = "    "
+INPUTDATA_URL = "https://osdf-data.gdex.ucar.edu/ncar/gdex/d651077/cesmdata/inputdata"
 
 
 def build_parser() -> argparse.ArgumentParser:
@@ -141,7 +144,9 @@ def normalize_paths(root: Path, relnames: Iterable[str]) -> List[Path]:
     return paths
 
 
-def stage_data(src: Path, inputdata_root: Path, staging_root: Path, check: bool = False) -> None:
+def stage_data(
+    src: Path, inputdata_root: Path, staging_root: Path, check: bool = False
+) -> None:
     """Stage a file by mirroring its path under `staging_root`.
 
     Destination path is computed by replacing the `inputdata_root` prefix of `src`
@@ -173,6 +178,10 @@ def stage_data(src: Path, inputdata_root: Path, staging_root: Path, check: bool
                 f"({staging_root})"
             )
         print(f"{INDENT}File is already published and linked.")
+        if can_file_be_downloaded(src.resolve(), staging_root):
+            print(f"{INDENT}File is available for download.")
+        else:
+            print(f"{INDENT}File is not (yet) available for download.")
         return
 
     # TODO: Check whether file is published but not relinked.
@@ -254,6 +263,38 @@ def get_staging_root() -> Path:
     return DEFAULT_STAGING_ROOT
 
 
+def can_file_be_downloaded(file_relpath: Path, staging_root: Path, timeout: float = 10):
+    """Check whether a file is available for download from the CESM inputdata server.
+
+    Sends a HEAD request to the CESM inputdata URL to verify if the file exists and is
+    accessible without downloading the entire file.
+
+    Args:
+        file_relpath: Relative path to the file (relative to staging_root), or an absolute
+            path that will be made relative to staging_root.
+        staging_root: Root directory of the staging area, used to compute relative path
+            if file_relpath is absolute.
+        timeout: Maximum time in seconds to wait for the server response. Default is 10.
+
+    Returns:
+        bool: True if the file is accessible (HTTP status 2xx or 3xx), False otherwise
+            (including 404, network errors, timeouts, etc.).
+    """
+    # Get URL
+    if file_relpath.is_absolute():
+        file_relpath = file_relpath.relative_to(staging_root)
+    url = os.path.join(INPUTDATA_URL, file_relpath)
+
+    # Check whether URL can be accessed
+    req = Request(url, method="HEAD")
+    try:
+        with urlopen(req, timeout=timeout) as resp:
+            return 200 <= resp.status < 400
+    except HTTPError:
+        # Server reached, but resource doesn't exist (404, 410, etc.)
+        return False
+
+
 def main(argv: List[str] | None = None) -> int:
     """Main entry point for the rimport tool.
 

tests/rimport/test_can_file_be_downloaded.py

@@ -0,0 +1,61 @@
+"""
+Tests for can_file_be_downloaded() function in rimport script.
+"""
+
+import os
+import sys
+import importlib.util
+from importlib.machinery import SourceFileLoader
+from pathlib import Path
+
+from shared import DEFAULT_STAGING_ROOT
+
+# Import rimport module from file without .py extension
+rimport_path = os.path.join(
+    os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
+    "rimport",
+)
+loader = SourceFileLoader("rimport", rimport_path)
+spec = importlib.util.spec_from_loader("rimport", loader)
+if spec is None:
+    raise ImportError(f"Could not create spec for rimport from {rimport_path}")
+rimport = importlib.util.module_from_spec(spec)
+sys.modules["rimport"] = rimport
+loader.exec_module(rimport)
+
+RELPATH_THAT_DOES_EXIST = os.path.join(
+    "share", "meshes", "ne3pg3_ESMFmesh_c221214_cdf5.asc"
+)
+
+
+class TestCanFileBeDownloaded:
+    """Test suite for can_file_be_downloaded() function."""
+    
+    def test_existing_file_exists(self):
+        """Test that the file that should exist does. If not, other tests will definitely fail."""
+        file_abspath = Path(os.path.join(DEFAULT_STAGING_ROOT, RELPATH_THAT_DOES_EXIST))
+        assert file_abspath.exists()
+
+    def test_true_abspath(self):
+        """Test that can_file_be_downloaded() is true for an existing file given absolute path"""
+        file_abspath = Path(os.path.join(DEFAULT_STAGING_ROOT, RELPATH_THAT_DOES_EXIST))
+        assert rimport.can_file_be_downloaded(
+            file_abspath,
+            DEFAULT_STAGING_ROOT,
+        )
+
+    def test_true_relpath(self):
+        """Test that can_file_be_downloaded() is true for an existing file given relative path"""
+        file_relpath = Path(RELPATH_THAT_DOES_EXIST)
+        assert rimport.can_file_be_downloaded(
+            file_relpath,
+            DEFAULT_STAGING_ROOT,
+        )
+
+    def test_false_nonexistent(self):
+        """Test that can_file_be_downloaded() is false for a nonexistent file"""
+        file_relpath = Path("weurueridniduafnea/smfnigsroerij/msdif8ernnr.nc")
+        assert not rimport.can_file_be_downloaded(
+            file_relpath,
+            DEFAULT_STAGING_ROOT,
+        )

tests/rimport/test_stage_data.py

@@ -88,12 +88,12 @@ def test_preserves_directory_structure(self, inputdata_root, staging_root):
         assert dst.exists()
         assert dst.read_text() == "nested data"
 
-    def test_prints_live_symlink_already_published(
+    def test_prints_live_symlink_already_published_not_downloadable(
         self, inputdata_root, staging_root, capsys
     ):
         """
         Test that staging a live, already-published symlink prints a message and returns
-        immediately without copying anything.
+        immediately without copying anything. Should say it's not available for download.
         """
         # Create a real file in staging and a symlink to it in inputdata
         real_file = staging_root / "real_file.nc"
@@ -106,10 +106,12 @@ def test_prints_live_symlink_already_published(
             # Should print message for live symlink and return early
             rimport.stage_data(src, inputdata_root, staging_root)
 
-            # Verify the right message was printed
+            # Verify the right messages were printed
             stdout = capsys.readouterr().out.strip()
             msg = "File is already published and linked"
             assert msg in stdout
+            msg = "File is not (yet) available for download"
+            assert msg in stdout
 
             # Verify the WRONG message was NOT printed
             msg = "is already under staging directory"
@@ -118,6 +120,40 @@ def test_prints_live_symlink_already_published(
             # Verify that shutil.copy2 was never called (function returned early)
             mock_copy.assert_not_called()
 
+    def test_prints_live_symlink_already_published_is_downloadable(
+        self, inputdata_root, staging_root, capsys
+    ):
+        """
+        Like test_prints_live_symlink_already_published_not_downloadable, but mocks
+        can_file_be_downloaded() to test "is available for download" message.
+        """
+        # Create a real file in staging and a symlink to it in inputdata
+        real_file = staging_root / "real_file.nc"
+        real_file.write_text("data")
+        src = inputdata_root / "link.nc"
+        src.symlink_to(real_file)
+
+        # Mock shutil.copy2 to verify it's never called
+        with patch("shutil.copy2") as mock_copy:
+            # Mock can_file_be_downloaded to return True
+            with patch("rimport.can_file_be_downloaded", return_value=True):
+                # Should print message for live symlink and return early
+                rimport.stage_data(src, inputdata_root, staging_root)
+
+                # Verify that shutil.copy2 was never called (function returned early)
+                mock_copy.assert_not_called()
+
+        # Verify the right messages were printed
+        stdout = capsys.readouterr().out.strip()
+        msg = "File is already published and linked"
+        assert msg in stdout
+        msg = "File is available for download"
+        assert msg in stdout
+
+        # Verify the WRONG message was NOT printed
+        msg = "is already under staging directory"
+        assert msg not in stdout
+
     def test_raises_error_for_live_symlink_pointing_somewhere_other_than_staging(
         self, tmp_path, inputdata_root, staging_root
     ):