Merge pull request #23 from DACCS-Climate/add-tests

Add tests

Author: mishaschwartz

.github/workflows/tests.yml

@@ -0,0 +1,21 @@
+name: Unit tests
+on:
+  pull_request:
+    types: [opened, synchronize, reopened, ready_for_review]
+jobs:
+  test:
+    if: github.event.pull_request.draft == false
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up python for testing
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.x"
+          cache: 'pip'
+      - name: Install python test dependencies
+        run: |
+          pip install .[test]
+      - name: Test with pytest
+        run: |
+          pytest ./tests/

README.md

@@ -21,14 +21,14 @@ The first thing to do is to get a `client` object:
 All the information about the network can now be retrieved from the `client` object. E.g. the nodes available in the network can be accessed as:
 ```python
 >>> client.nodes
-{'UofT': <marble_client.node.MarbleNode at 0x10c129990>,
- 'PAVICS': <marble_client.node.MarbleNode at 0x10c6dd690>,
- 'Hirondelle': <marble_client.node.MarbleNode at 0x10c6dd890>}
+{'UofTRedOak': <MarbleNode(id: 'UofTRedOak', name: 'Red Oak')>, 
+ 'PAVICS': <MarbleNode(id: 'PAVICS', name: 'PAVICS')>, 
+ 'Hirondelle': <MarbleNode(id: 'Hirondelle', name: 'Hirondelle')>}
 ```
 The returned object is a python `dict` with node names for keys and `MarbleNode` objects as values. A particular node can be accessed as:
 
 ```python
->>> mynode = client['UofT']
+>>> mynode = client['UofTRedOak']
 >>> type(mynode)
 marble_client.node.MarbleNode
 ```
@@ -43,7 +43,7 @@ True
 The URL for the node can be retrieved as:
 ```python
 >>> mynode.url
-'https://daccs.cs.toronto.edu'
+'https://redoak.cs.toronto.edu'
 ```
 
 Various other qualities about the node can be accessed as shown below (see [implementation](https://github.com/DACCS-Climate/marble_client_python/blob/main/marble_client/node.py) for the full list of available attributes).
@@ -121,7 +121,7 @@ Get the node your notebook/script is currently running on:
 ```python
 >>> client = MarbleClient()
 >>> client.this_node
-<marble_client.node.MarbleNode at 0x10c129990>
+<MarbleNode(id: 'UofTRedOak', name: 'Red Oak')>
 ```
 
 Add session cookies to a `requests.Session` object. This means that any request made with that session variable will

marble_client/__init__.py

@@ -1,4 +1,4 @@
 from .client import MarbleClient
-from .exceptions import MarbleBaseError, ServiceNotAvailableError, UnknownNodeError
+from .exceptions import MarbleBaseError, ServiceNotAvailableError, UnknownNodeError, JupyterEnvironmentError
 from .node import MarbleNode
 from .services import MarbleService

marble_client/client.py

@@ -1,17 +1,16 @@
-import contextlib
 import datetime
 import json
 import os
 import shutil
 import warnings
 from functools import wraps, cache
-from typing import Optional
+from typing import Optional, Any
 from urllib.parse import urlparse
 
 import dateutil.parser
 import requests
 
-from marble_client.constants import CACHE_FNAME, CACHE_META_FNAME, NODE_REGISTRY_URL
+from marble_client.constants import CACHE_FNAME, NODE_REGISTRY_URL
 from marble_client.exceptions import UnknownNodeError, JupyterEnvironmentError
 from marble_client.node import MarbleNode
 
@@ -25,45 +24,43 @@ def check_jupyterlab(f):
 
     This is used as a pre-check for functions that only work in a Marble Jupyterlab
     environment.
+
+    Note that this checks if either the BIRDHOUSE_HOST_URL or PAVICS_HOST_URL are present to support
+    versions of birdhouse-deploy prior to 2.4.0.
     """
     @wraps(f)
     def wrapper(*args, **kwargs):
-        if os.getenv("PAVICS_HOST_URL"):
+        birdhouse_host_var = ("PAVICS_HOST_URL", "BIRDHOUSE_HOST_URL")
+        jupyterhub_env_vars = ("JUPYTERHUB_API_URL", "JUPYTERHUB_USER", "JUPYTERHUB_API_TOKEN")
+        if any(os.getenv(var) for var in birdhouse_host_var) and all(os.getenv(var) for var in jupyterhub_env_vars):
             return f(*args, **kwargs)
         raise JupyterEnvironmentError("Not in a Marble jupyterlab environment")
     return wrapper
 
 
 class MarbleClient:
-    def __init__(self, fallback: Optional[bool] = True) -> None:
+    _registry_cache_key = "marble_client_python:cached_registry"
+    _registry_cache_last_updated_key = "marble_client_python:last_updated"
+
+    def __init__(self, fallback: bool = True) -> None:
         """Constructor method
 
         :param fallback: If True, then fall back to a cached version of the registry
             if the cloud registry cannot be accessed, defaults to True
-        :type fallback: Optional[bool], optional
+        :type fallback: bool
         :raises requests.exceptions.RequestException: Raised when there is an issue
             connecting to the cloud registry and `fallback` is False
         :raises UserWarning: Raised when there is an issue connecting to the cloud registry
             and `fallback` is True
         :raise RuntimeError: If cached registry needs to be read but there is no cache
         """
-        self._fallback = fallback
         self._nodes: dict[str, MarbleNode] = {}
-        self._registry: dict = {}
-        try:
-            registry = requests.get(NODE_REGISTRY_URL)
-            registry.raise_for_status()
-        except (requests.exceptions.RequestException, requests.exceptions.ConnectionError):
-            if self._fallback:
-                warnings.warn("Cannot retrieve cloud registry. Falling back to cached version")
-                self._load_registry_from_cache()
-            else:
-                raise
-        else:
-            self._load_registry_from_cloud(registry)
+        self._registry_uri: str
+        self._registry: dict
+        self._registry_uri, self._registry = self._load_registry(fallback)
 
-        for node, node_details in self._registry.items():
-            self._nodes[node] = MarbleNode(node, node_details)
+        for node_id, node_details in self._registry.items():
+            self._nodes[node_id] = MarbleNode(node_id, node_details)
 
     @property
     def nodes(self) -> dict[str, MarbleNode]:
@@ -78,11 +75,13 @@ def this_node(self) -> MarbleNode:
 
         Note that this function only works in a Marble Jupyterlab environment.
         """
-        host_url = urlparse(os.getenv("PAVICS_HOST_URL"))
+        # PAVICS_HOST_URL is the deprecated variable used in older versions (<2.4.0) of birdhouse-deploy
+        url_string = os.getenv("BIRDHOUSE_HOST_URL", os.getenv("PAVICS_HOST_URL"))
+        host_url = urlparse(url_string)
         for node in self.nodes.values():
             if urlparse(node.url).hostname == host_url.hostname:
                 return node
-        raise UnknownNodeError(f"No node found in the registry with the url {host_url}")
+        raise UnknownNodeError(f"No node found in the registry with the url '{url_string}'")
 
     @check_jupyterlab
     def this_session(self, session: Optional[requests.Session] = None) -> requests.Session:
@@ -97,78 +96,94 @@ def this_session(self, session: Optional[requests.Session] = None) -> requests.S
             session = requests.Session()
         r = requests.get(f"{os.getenv('JUPYTERHUB_API_URL')}/users/{os.getenv('JUPYTERHUB_USER')}",
                          headers={"Authorization": f"token {os.getenv('JUPYTERHUB_API_TOKEN')}"})
+        try:
+            r.raise_for_status()
+        except requests.HTTPError as err:
+            raise JupyterEnvironmentError("Cannot retrieve login cookies through the JupyterHub API.") from err
         for name, value in r.json().get("auth_state", {}).get("magpie_cookies", {}).items():
             session.cookies.set(name, value)
         return session
 
+    @property
+    def registry_uri(self):
+        return self._registry_uri
+
     def __getitem__(self, node: str) -> MarbleNode:
         try:
             return self.nodes[node]
-        except KeyError:
-            raise UnknownNodeError(f"No node named '{node}' in the Marble network.") from None
+        except KeyError as err:
+            raise UnknownNodeError(f"No node named '{node}' in the Marble network.") from err
 
-    def _load_registry_from_cloud(self, registry_response: requests.Response) -> None:
-        try:
-            self._registry = registry_response.json()
-        except json.JSONDecodeError:
-            raise RuntimeError(
-                "Could not parse JSON returned from the cloud registry. "
-                f"Consider re-trying with 'fallback' set to True when instantiating the {self.__class__.__name__} "
-                "object."
-            )
-        self._save_registry_as_cache()
-
-    def _load_registry_from_cache(self):
-        try:
-            with open(CACHE_FNAME, "r") as f:
-                self._registry = json.load(f)
-        except FileNotFoundError:
-            raise RuntimeError(f"Local registry cache not found. No file named {CACHE_FNAME}.") from None
+    def __contains__(self, node: str) -> bool:
+        """Check if a node is available
+
+        :param node: ID of the Marble node
+        :type node: str
+        :return: True if the node is present in the registry, False otherwise
+        :rtype: bool
+        """
+        return node in self.nodes
 
+    def _load_registry(self, fallback: bool = True) -> tuple[str, dict[str, Any]]:
         try:
-            with open(CACHE_META_FNAME, "r") as f:
-                data = json.load(f)
-                date = dateutil.parser.isoparse(data["last_cache_date"])
-        except (FileNotFoundError, ValueError):
-            date = "Unknown"
+            registry_response = requests.get(NODE_REGISTRY_URL)
+            registry_response.raise_for_status()
+            registry = registry_response.json()
+        except (requests.exceptions.RequestException, requests.exceptions.ConnectionError) as err:
+            error = err
+            error_msg = f"Cannot retrieve registry from {NODE_REGISTRY_URL}."
+        except json.JSONDecodeError as err:
+            error = err
+            error_msg = f"Could not parse JSON returned from the registry at {NODE_REGISTRY_URL}"
+        else:
+            self._save_registry_as_cache(registry)
+            return NODE_REGISTRY_URL, registry
 
-        print(f"Registry loaded from cache dating: {date}")
+        if fallback:
+            warnings.warn(f"{error_msg} Falling back to cached version")
+            return f"file://{os.path.realpath(CACHE_FNAME)}", self._load_registry_from_cache()
+        else:
+            raise RuntimeError(error_msg) from error
 
-    def _save_registry_as_cache(self):
+    def _load_registry_from_cache(self) -> dict[str, Any]:
+        try:
+            with open(CACHE_FNAME) as f:
+                cached_registry = json.load(f)
+        except FileNotFoundError as err:
+            raise RuntimeError(f"Local registry cache not found. No file named {CACHE_FNAME}.") from err
+        except json.JSONDecodeError as err:
+            raise RuntimeError(f"Could not parse JSON returned from the cached registry at {CACHE_FNAME}") from err
+        else:
+            if self._registry_cache_key in cached_registry:
+                registry = cached_registry[self._registry_cache_key]
+                date = dateutil.parser.isoparse(cached_registry[self._registry_cache_last_updated_key])
+            else:
+                # registry is cached in old format, re-cache it in the newer format
+                registry = cached_registry
+                self._save_registry_as_cache(registry)
+                date = "Unknown"
+            print(f"Registry loaded from cache dating: {date}")
+            return registry
+
+    def _save_registry_as_cache(self, registry: dict[str, Any]) -> None:
         cache_backup = CACHE_FNAME + ".backup"
-        cache_meta_backup = CACHE_META_FNAME + ".backup"
 
         # Create cache parent directories if they don't exist
-        for cache_dir in {os.path.dirname(CACHE_FNAME), os.path.dirname(CACHE_META_FNAME)}:
-            os.makedirs(cache_dir, exist_ok=True)
-
-        # Suppressing a FileNotFoundError error for the first use case where a cached registry file
-        # does not exist
-        with contextlib.suppress(FileNotFoundError):
+        os.makedirs(os.path.dirname(CACHE_FNAME), exist_ok=True)
+        if os.path.isfile(CACHE_FNAME):
             shutil.copy(CACHE_FNAME, cache_backup)
-            shutil.copy(CACHE_META_FNAME, cache_meta_backup)
 
         try:
-            metadata = {"last_cache_date": datetime.datetime.now(tz=datetime.timezone.utc).isoformat()}
-
-            # Write the metadata
-            with open(CACHE_META_FNAME, "w") as f:
-                json.dump(metadata, f)
-
-            # write the registry
             with open(CACHE_FNAME, "w") as f:
-                json.dump(self._registry, f)
-
-        except OSError as e:
-            # If either the cache file or the metadata file could not be written, then restore from backup files
+                data = {self._registry_cache_key: registry,
+                        self._registry_cache_last_updated_key: datetime.datetime.now(tz=datetime.timezone.utc).isoformat()}
+                json.dump(data, f)
+        except OSError:
+            # If the cache file cannot be written, then restore from backup files
             shutil.copy(cache_backup, CACHE_FNAME)
-            shutil.copy(cache_meta_backup, CACHE_META_FNAME)
-
         finally:
-            # Similarly, suppressing an error that I don't need to catch
-            with contextlib.suppress(FileNotFoundError):
+            if os.path.isfile(cache_backup):
                 os.remove(cache_backup)
-                os.remove(cache_meta_backup)
 
 
 if __name__ == "__main__":

marble_client/constants.py

@@ -2,7 +2,7 @@
 
 from platformdirs import user_cache_dir
 
-__all__ = ("NODE_REGISTRY_URL", "CACHE_FNAME", "CACHE_META_FNAME")
+__all__ = ("NODE_REGISTRY_URL", "CACHE_FNAME")
 
 # Marble node registry URL
 NODE_REGISTRY_URL: str = os.getenv(
@@ -14,6 +14,3 @@
 
 # location to write registry cache
 CACHE_FNAME: str = os.path.join(_CACHE_DIR, "registry.cached.json")
-
-# location to write metadata about the registry cache
-CACHE_META_FNAME: str = os.path.join(_CACHE_DIR, "cache_metadata.json")