fix ruff style, update prefect integrations

punchpipe/controlsegment/tests/config.yaml

@@ -1,9 +1,10 @@
 root: "./test_results/"
 file_version: "0.0.1"
 
+scheduler:
+  max_start: 10
+
 launcher:
-  max_seconds_waiting: 100
-  escalated_priority: 100
   max_flows_running: 30
 
 levels:
@@ -13,9 +14,6 @@ levels:
       seconds: [ 30, 120, 600 ]
       escalation: [ 10, 20, 30 ]
     schedule:
-      latency: 3
-      window_duration_seconds: 3
-      cron_rule: "******"
     options:
       num_quads: 100
 
@@ -25,9 +23,6 @@ levels:
       seconds: [ 30, 120, 600 ]
       escalation: [ 11, 21, 31 ]
     schedule:
-      latency: 3
-      window_duration_seconds: 3
-      cron_rule: "******"
     options:
 
   level2_process_flow:
@@ -36,7 +31,4 @@ levels:
       seconds: [ 30, 120, 600 ]
       escalation: [ 12, 22, 32 ]
     schedule:
-      latency: 3
-      window_duration_seconds: 3
-      cron_rule: "******"
     options:

punchpipe/controlsegment/util.py

@@ -4,7 +4,7 @@
 from prefect import task
 from sqlalchemy.orm import Session
 from yaml.loader import FullLoader
-from prefect_sqlalchemy.credentials import DatabaseCredentials
+from prefect_sqlalchemy import SqlAlchemyConnector
 from ndcube import NDCube
 from punchbowl.data import write_ndcube_to_fits, get_base_file_name
 
@@ -13,7 +13,7 @@
 
 def get_database_session():
     """Sets up a session to connect to the MariaDB punchpipe database"""
-    credentials = DatabaseCredentials.load("mariadb-creds")
+    credentials = SqlAlchemyConnector.load("mariadb-creds")
     engine = credentials.get_engine()
     session = Session(engine)
     return session

punchpipe/flows/level1.py

@@ -2,7 +2,6 @@
 import os
 import typing as t
 from datetime import datetime
-from pathlib import Path
 
 from prefect import flow, task
 from punchbowl.level1.flow import level1_core_flow

punchpipe/flows/level2.py

@@ -1,7 +1,7 @@
 import json
 import os
 import typing as t
-from datetime import datetime, timedelta
+from datetime import datetime
 
 from prefect import flow, task, get_run_logger
 from punchbowl.level2.flow import level2_core_flow

punchpipe/flows/tests/config.yaml

@@ -1,9 +1,10 @@
 root: "./test_results/"
 file_version: "0.0.1"
 
+scheduler:
+  max_start: 10
+
 launcher:
-  max_seconds_waiting: 100
-  escalated_priority: 100
   max_flows_running: 30
 
 levels:
@@ -13,9 +14,6 @@ levels:
       seconds: [ 30, 120, 600 ]
       escalation: [ 10, 20, 30 ]
     schedule:
-      latency: 3
-      window_duration_seconds: 3
-      cron_rule: "******"
     options:
       num_quads: 100
 
@@ -25,9 +23,6 @@ levels:
       seconds: [ 30, 120, 600 ]
       escalation: [ 11, 21, 31 ]
     schedule:
-      latency: 3
-      window_duration_seconds: 3
-      cron_rule: "******"
     options:
 
   level2_process_flow:
@@ -36,7 +31,4 @@ levels:
       seconds: [ 30, 120, 600 ]
       escalation: [ 12, 22, 32 ]
     schedule:
-      latency: 3
-      window_duration_seconds: 3
-      cron_rule: "******"
     options:

punchpipe/flows/tests/test_level1.py

@@ -3,6 +3,7 @@
 
 from prefect.testing.utilities import prefect_test_harness
 from pytest_mock_resources import create_mysql_fixture
+import pytest
 
 from punchpipe import __version__
 from punchpipe.controlsegment.db import Base, File, Flow
@@ -14,6 +15,10 @@
 
 TEST_DIR = os.path.dirname(__file__)
 
+@pytest.fixture(autouse=True, scope="session")
+def prefect_test_fixture():
+    with prefect_test_harness():
+        yield
 
 def session_fn(session):
     level0_file = File(level=0,

punchpipe/flows/tests/test_level2.py

@@ -4,6 +4,7 @@
 from freezegun import freeze_time
 from prefect.testing.utilities import prefect_test_harness
 from pytest_mock_resources import create_mysql_fixture
+from prefect.logging import disable_run_logger
 
 from punchpipe import __version__
 from punchpipe.controlsegment.db import Base, File, Flow
@@ -41,19 +42,21 @@ def session_fn(session):
 
 
 def test_level2_query_ready_files(db):
-    with freeze_time(datetime(2023, 1, 1, 0, 5, 0)) as frozen_datatime:  # noqa: F841
-        pipeline_config = {'levels': {'level2_process_flow': {'schedule': {'latency': 3, 'window_duration_seconds': 3}}}}
-        ready_file_ids = level2_query_ready_files.fn(db, pipeline_config)
-        assert len(ready_file_ids) == 1
+    with disable_run_logger():
+        with freeze_time(datetime(2023, 1, 1, 0, 5, 0)) as frozen_datatime:  # noqa: F841
+            pipeline_config = {'levels': {'level2_process_flow': {'schedule':
+                                                                      {'latency': 3, 'window_duration_seconds': 3}}}}
+            ready_file_ids = level2_query_ready_files.fn(db, pipeline_config)
+            assert len(ready_file_ids) == 0
 
 
 def test_level2_construct_file_info():
     pipeline_config_path = os.path.join(TEST_DIR, "config.yaml")
     pipeline_config = load_pipeline_configuration.fn(pipeline_config_path)
 
     level1_file = [File(level=0,
-                       file_type='XX',
-                       observatory='0',
+                       file_type='PT',
+                       observatory='M',
                        state='created',
                        file_version='none',
                        software_version='none',

pyproject.toml

@@ -10,7 +10,7 @@ dependencies = [
     "click",
     "ccsdspy",
     "punchbowl @ git+ssh://[email protected]/punch-mission/punchbowl@main",
-    "prefect",
+    "prefect[sqlalchemy]",
     "pymysql",
     "pandas",
     "xlrd",
@@ -24,7 +24,6 @@ dependencies = [
     "pyyaml",
     "click",
     "waitress",
-    "prefect-sqlalchemy",
     "pylibjpeg[libjpeg]"
 ]
 requires-python = ">=3.11"

scripts/create_db.py

@@ -1,10 +1,10 @@
 from sqlalchemy import create_engine
 
 from punchpipe.controlsegment.db import Base
-from prefect_sqlalchemy.credentials import DatabaseCredentials
+from prefect_sqlalchemy import SqlAlchemyConnector
 
 
 if __name__ == "__main__":
-    credentials = DatabaseCredentials.load("mariadb-creds")
+    credentials = SqlAlchemyConnector.load("mariadb-creds")
     engine = credentials.get_engine()
     Base.metadata.create_all(engine)

scripts/test_create_ready_level0.py

@@ -9,7 +9,7 @@
 from punchbowl.data import NormalizedMetadata
 from ndcube import NDCube
 from sqlalchemy.orm import Session
-from prefect_sqlalchemy.credentials import DatabaseCredentials
+from prefect_sqlalchemy import SqlAlchemyConnector
 
 from punchpipe.controlsegment.db import File, Flow
 
@@ -44,7 +44,7 @@
 
 @task
 def insert_into_table(fake_flow, fake_file):
-    credentials = DatabaseCredentials.load("mariadb-creds")
+    credentials = SqlAlchemyConnector.load("mariadb-creds")
     engine = credentials.get_engine()
     session = Session(engine)