Add tests

.github/workflows/daily-integration-tests.yml

@@ -0,0 +1,32 @@
+name: Daily integration tests
+
+on:
+  schedule:
+    - cron: "17 6 * * *"
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  integration:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1
+        with:
+          fetch-depth: 0
+          persist-credentials: false
+
+      - name: Install Pixi environment
+        uses: prefix-dev/setup-pixi@d3f436a425481402e6a95a1d1fc10331c708cd9e
+        with:
+          pixi-version: v0.77.0
+          manifest-path: pyproject.toml
+          environments: dev
+          locked: true
+          cache: true
+
+      - name: Run integration tests
+        run: pixi run -e dev test-integration

.github/workflows/unit-tests.yml

@@ -0,0 +1,33 @@
+name: Unit tests
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+permissions:
+  contents: read
+
+jobs:
+  unit:
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1
+        with:
+          fetch-depth: 0
+          persist-credentials: false
+
+      - name: Install Pixi environment
+        uses: prefix-dev/setup-pixi@d3f436a425481402e6a95a1d1fc10331c708cd9e
+        with:
+          pixi-version: v0.77.0
+          manifest-path: pyproject.toml
+          environments: dev
+          locked: true
+          cache: true
+
+      - name: Run unit tests
+        run: pixi run -e dev test-unit

README.md

@@ -1,3 +1,5 @@
 # tst-sim-tools
 
+[![Unit tests](https://github.com/NSLS2/tst-sim-tools/actions/workflows/unit-tests.yml/badge.svg)](https://github.com/NSLS2/tst-sim-tools/actions/workflows/unit-tests.yml)
+
 Tools for the TST beamline's simulated "endstation"

pyproject.toml

@@ -36,6 +36,7 @@ dependencies = [
 
 [dependency-groups]
 dev = [
+  "h5py",
   "ophyd-async[ca,sim]",
   "import-linter",
   "numpydoc",
@@ -46,6 +47,7 @@ dev = [
   "pytest-asyncio",
   "pytest-cov",
   "pytest-mock",
+  "pytest-socket>=0.8.1",
   "ruff",
   "tiled[client,server]>=0.2.14",
   "bluesky-tiled-plugins>=2.0.9",
@@ -79,6 +81,13 @@ allow-direct-references = true
 [tool.pyright]
 typeCheckingMode = "standard"
 
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "--strict-config --strict-markers --disable-socket --allow-unix-socket"
+markers = [
+  "integration: local round-trip tests excluded from per-push CI",
+]
+
 [tool.coverage.paths]
 # Tests are run from installed location, map back to the src directory
 source = ["src"]
@@ -198,6 +207,10 @@ qs = { features = [
 dev = { features = ["dev"], solve-group = "default" }
 docs = { features = ["docs"], solve-group = "default" }
 
+[tool.pixi.feature.dev.tasks]
+test-unit = "pytest tests/unit --cov=tst_sim_tools --cov-report=term-missing"
+test-integration = "pytest -m integration tests/integration"
+
 [tool.pixi.feature.startup-local.tasks]
 start-local = "TILED=0 ipython -i -m tst_sim_tools.startup"
 start-staging = "TILED=1 ipython -i -m tst_sim_tools.startup"

src/tst_sim_tools/agents/energy_alignment.py

@@ -1,7 +1,7 @@
 """Blop agent for energy dependent alignment."""
 
 import time
-from collections.abc import Sequence
+from collections.abc import Hashable, Mapping, Sequence
 from functools import partial
 from typing import Any, cast
 
@@ -31,6 +31,8 @@
 MIN_INTENSITY = "min_intensity"
 DEFAULT_IMAGE_THRESHOLD = 0.02
 DEFAULT_IMAGE_BLUR = 1.0
+CATALOG_POLL_INTERVAL = 0.1
+CATALOG_POLL_TIMEOUT = 10.0
 BMM_ENERGY_ALIGNMENT_MAX_CENTROID_ERROR_WEIGHT = 0.25
 BMM_ENERGY_ALIGNMENT_CENTROID_SPAN_WEIGHT = 0.1
 BMM_ENERGY_ALIGNMENT_FWHM_WEIGHT = 0.005
@@ -73,19 +75,29 @@ def __init__(
         self._threshold = threshold
         self._blur = blur
 
-    def _poll_for_images(self, uid: str) -> np.ndarray:
+    def _poll_for_run_images(self, uid: Hashable) -> tuple[Any, np.ndarray]:
+        deadline = time.monotonic() + CATALOG_POLL_TIMEOUT
         while True:
+            waiting_for = f"run {uid!r} in Tiled"
             try:
                 run = self._client[uid]
+                waiting_for = f"primary stream for run {uid!r}"
                 stream = run["primary"]
-                return stream[self._image_key].read()
-            except KeyError:
-                time.sleep(0.1)
+                waiting_for = f"image key {self._image_key!r} in the primary stream for run {uid!r}"
+                image = stream[self._image_key]
+                waiting_for = f"readable data for image key {self._image_key!r} in run {uid!r}"
+                return run, image.read()
+            except KeyError as error:
+                if time.monotonic() >= deadline:
+                    raise TimeoutError(
+                        f"Timed out after {CATALOG_POLL_TIMEOUT:g} seconds waiting for {waiting_for}"
+                    ) from error
+                time.sleep(CATALOG_POLL_INTERVAL)
 
-    def __call__(self, uid: str, suggestions: list[dict]) -> list[dict]:
-        images = self._poll_for_images(uid)
+    def __call__(self, uid: Hashable, suggestions: Sequence[Mapping]) -> list[dict]:
+        run, images = self._poll_for_run_images(uid)
         image_stack = image_series(images)
-        suggestion_ids = [suggestion["_id"] for suggestion in self._client[uid].metadata["start"]["blop_suggestions"]]
+        suggestion_ids = [suggestion["_id"] for suggestion in run.metadata["start"]["blop_suggestions"]]
         n_energies = self._energies.size
         expected_images = len(suggestion_ids) * n_energies
         if image_stack.shape[0] != expected_images:

src/tst_sim_tools/analysis/image.py

@@ -7,6 +7,13 @@
 from scipy import ndimage, signal
 
 
+def _finite_float_array(values: ArrayLike) -> np.ndarray:
+    """Return an owned floating-point array with non-finite values set to zero."""
+    array = np.array(values, dtype=float, copy=True)
+    array[~np.isfinite(array)] = 0.0
+    return array
+
+
 def image_series(images: ArrayLike) -> np.ndarray:
     """Return detector data as an ``(N, height, width)`` floating-point image stack.
 
@@ -21,7 +28,7 @@ def image_series(images: ArrayLike) -> np.ndarray:
     numpy.ndarray
         Floating-point image stack with non-finite values replaced by zero.
     """
-    stack = np.asarray(images, dtype=float)
+    stack = _finite_float_array(images)
     if stack.ndim == 2:
         stack = stack[np.newaxis, :, :]
     elif stack.ndim == 4:
@@ -85,7 +92,7 @@ def threshold_image(image: ArrayLike, threshold: float = 0.0) -> np.ndarray:
     if not 0.0 <= threshold <= 1.0:
         raise ValueError(f"Expected threshold in [0, 1], but got {threshold}")
 
-    thresholded = np.asarray(image, dtype="float")
+    thresholded = _finite_float_array(image)
     peak = float(thresholded.max())
     if peak <= 0.0:
         thresholded.fill(0.0)
@@ -118,7 +125,7 @@ def gaussian_blur(image: ArrayLike, sigma: float = 0.0, truncate: float = 4.0) -
     if truncate <= 0.0:
         raise ValueError(f"Expected positive truncate, but got {truncate}")
 
-    blurred = np.asarray(image, dtype="float")
+    blurred = _finite_float_array(image)
     if sigma == 0.0:
         return blurred
 
@@ -275,8 +282,8 @@ def analyze_image(
     """
     processed = preprocess(image, threshold=threshold, blur=blur)
     height, width = processed.shape
-    center_x = width
-    center_y = height
+    center_x = (width - 1) / 2.0
+    center_y = (height - 1) / 2.0
     total = float(processed.sum())
     peak = float(processed.max())
     if total <= 0.0:

src/tst_sim_tools/startup.py

@@ -13,7 +13,7 @@
 from bluesky.plan_stubs import mv, rd
 from bluesky.plans import count, grid_scan, list_scan, rel_scan, scan
 from bluesky_tiled_plugins import TiledWriter
-from nslsii.ophyd_async.providers import NSLS2PathProvider
+from nslsii.ophyd_async.providers import NSLS2PathProvider  # pyright: ignore[reportMissingImports]
 from ophyd_async.core import UUIDFilenameProvider, YMDPathProvider, init_devices
 
 from tst_sim_tools.devices.detectors import XRTScreenDetector

tests/conftest.py

@@ -0,0 +1,48 @@
+import asyncio
+import os
+from collections import defaultdict
+from collections.abc import Generator
+
+os.environ.setdefault("MPLBACKEND", "Agg")
+
+import matplotlib.pyplot as plt
+import pytest
+from bluesky import RunEngine
+
+
+@pytest.fixture(autouse=True)
+def close_figures() -> Generator[None]:
+    """Close every Matplotlib figure after each test."""
+    yield
+    plt.close("all")
+
+
+@pytest.fixture
+def run_engine() -> Generator[RunEngine]:
+    """Provide a RunEngine with an isolated, deterministically closed loop."""
+    loop = asyncio.new_event_loop()
+    engine = RunEngine({}, call_returns_result=True, loop=loop)
+    try:
+        yield engine
+    finally:
+        if engine.state not in ("idle", "panicked"):
+            try:
+                engine.halt()
+            except RuntimeError:
+                pass
+        loop.call_soon_threadsafe(loop.stop)
+        thread = getattr(engine, "_th", None)
+        if thread is not None:
+            thread.join()
+        loop.close()
+
+
+@pytest.fixture
+def documents():
+    """Collect emitted Bluesky documents by name."""
+    collected = defaultdict(list)
+
+    def collect(name, document):
+        collected[name].append(document)
+
+    return collected, collect