diff --git a/punchpipe/controlsegment/util.py b/punchpipe/controlsegment/util.py index 58d7b50..643f1d8 100644 --- a/punchpipe/controlsegment/util.py +++ b/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 diff --git a/punchpipe/flows/tests/test_level2.py b/punchpipe/flows/tests/test_level2.py index 20f9c0e..d9ebe77 100644 --- a/punchpipe/flows/tests/test_level2.py +++ b/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,10 +42,12 @@ 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(): @@ -52,8 +55,8 @@ def test_level2_construct_file_info(): 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', diff --git a/pyproject.toml b/pyproject.toml index 7ef97df..26e3e57 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ dependencies = [ "click", "ccsdspy", "punchbowl @ git+ssh://git@github.com/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"