Add testing

.github/workflows/ci.yml

@@ -0,0 +1,41 @@
+# This workflow will install Python dependencies, run tests and lint with a variety of Python versions
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python
+
+name: Python package
+
+on:
+  push:
+    branches: ["main"]
+  pull_request:
+    branches: ["main"]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.9", "3.10", "3.11"]
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v3
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install flake8 pytest prefect tiled bluesky-tiled-plugins
+          if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+      - name: Lint with flake8
+        run: |
+          # stop the build if there are Python syntax errors or undefined names
+          flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
+          # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+          flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+      - name: Test with pytest
+        env:
+          TEST_TILED_API_KEY: ${{ secrets.TEST_TILED_API_KEY }}
+        run: |
+          pytest

data_validation.py

@@ -1,28 +1,32 @@
 from prefect import task, flow, get_run_logger
-from prefect.blocks.system import Secret
 import time as ttime
-from tiled.client import from_profile
+from tiled.client import from_uri
 
 
 @task(retries=2, retry_delay_seconds=10)
-def read_all_streams(uid, beamline_acronym):
+def get_run(uid, api_key=None):
+    cl = from_uri("https://tiled.nsls2.bnl.gov", api_key=api_key)
+    run = cl["tst/raw"][uid]
+    return run
+
+
+@task(retries=2, retry_delay_seconds=10)
+def read_stream(run, stream):
+    return run[stream].read()
+
+
+@flow
+def data_validation(uid, api_key=None):
     logger = get_run_logger()
-    api_key = Secret.load("tiled-tst-api-key").get()
-    cl = from_profile("nsls2", api_key=api_key)
-    run = cl["tst"]["raw"][uid]
+    run = get_run(uid, api_key=api_key)
     logger.info(f"Validating uid {run.start['uid']}")
     start_time = ttime.monotonic()
     for stream in run:
         logger.info(f"{stream}:")
         stream_start_time = ttime.monotonic()
-        stream_data = run[stream].read()
+        stream_data = read_stream(run, stream)  # noqa: F841
         stream_elapsed_time = ttime.monotonic() - stream_start_time
         logger.info(f"{stream} elapsed_time = {stream_elapsed_time}")
         logger.info(f"{stream} nbytes = {stream_data.nbytes:_}")
     elapsed_time = ttime.monotonic() - start_time
     logger.info(f"{elapsed_time = }")
-
-
-@flow
-def data_validation(uid):
-    read_all_streams(uid, beamline_acronym="tst")

end_of_run_workflow.py

@@ -1,20 +1,34 @@
+import os
+
 from prefect import task, flow, get_run_logger
 from data_validation import data_validation
 from test_extra_client import get_other_docs
+from dotenv import load_dotenv
 # from long_flow import long_flow
 
 
 @task
-def log_completion():
+def get_api_key_from_env(api_key=None):
+    with open("/srv/container.secret", "r") as secrets:
+        load_dotenv(stream=secrets)
+    api_key = os.environ["TILED_API_KEY"]
+    return api_key
+
+
+@task
+def log_completion(dry_run=False):
     logger = get_run_logger()
-    logger.info("Complete")
+    logger.info(f"Complete! dry_run: {dry_run}")
 
 
 @flow
-def end_of_run_workflow(stop_doc):
+def end_of_run_workflow(stop_doc, dry_run=False, api_key=None):
     uid = stop_doc["run_start"]
     # hello_world()
-    data_validation(uid, return_state=True)
-    get_other_docs(uid)
-    # long_flow(iterations=100, sleep_length=10)
-    log_completion()
+    if not api_key:
+        api_key = get_api_key_from_env(api_key=api_key)
+    data_validation(uid, api_key=api_key)
+    get_other_docs(uid, api_key=api_key)
+    # long_flow(iterations=100, sleep_length=10, dry_run=dry_run)
+    log_completion(dry_run=dry_run)
+    return True

pixi.toml

@@ -8,6 +8,13 @@ prefect = "3.*"
 python = "<3.14"
 tiled-client = ">=0.2.3"
 bluesky-tiled-plugins = ">=2"
+python-dotenv = ">=1.2.1,<2"
 
 [pypi-dependencies]
 lixtools = "==2023.1.23.0"
+
+[feature.test.dependencies]
+pytest = "*"
+
+[environments]
+test = ["test"]

prefect.yaml

@@ -24,14 +24,12 @@ deployments:
     schedule: {}
     work_pool:
       job_variables:
-        env:
-          TILED_SITE_PROFILES: /nsls2/software/etc/tiled/profiles
         image: ghcr.io/nsls2/tst-workflows:main
         image_pull_policy: Always
         network_mode: slirp4netns
         userns: "keep-id:uid=402974,gid=402974" # workflow-tst:workflow-tst
         volumes:
           - /nsls2/data/tst/proposals:/nsls2/data/tst/proposals
-          - /nsls2/software/etc/tiled:/nsls2/software/etc/tiled
+          - /srv/prefect3-docker-worker-tst/app:/srv
         auto_remove: true
       name: tst-work-pool-docker

test_extra_client.py

@@ -1,10 +1,10 @@
 from prefect import task, get_run_logger
-from utils import get_tiled_client
+from data_validation import get_run
 
 
 @task
-def get_other_docs(uid):
+def get_other_docs(uid, api_key=None):
     logger = get_run_logger()
-    result = get_tiled_client()["raw"][uid]
+    result = get_run(uid, api_key=api_key)
     for name, doc in result.documents():
         logger.info(f"name: {name}, doc: {doc}")

test_run.py

@@ -0,0 +1,19 @@
+from prefect.logging import disable_run_logger
+from end_of_run_workflow import end_of_run_workflow
+import os
+import pytest
+
+
+@pytest.fixture(autouse=True, scope="session")
+def prefect_disable_logging():
+    with disable_run_logger():
+        yield
+
+
+def test_end_of_run_workflow(prefect_disable_logging):
+    print("starting test!")
+    assert end_of_run_workflow(
+        stop_doc={"run_start": "f0954c84-f652-4350-9f6d-44b724f4ed9f"},
+        api_key=os.environ["TEST_TILED_API_KEY"],
+    )
+    print("finished test!")