minor cleanup

Author: timbeccue

banzai/cache/replication.py

@@ -1,5 +1,4 @@
 from sqlalchemy import text, create_engine
-from banzai import dbs
 from banzai.logs import get_logger
 
 # PostgreSQL logical replication is managed via server-level DDL commands

banzai/dbs.py

@@ -14,7 +14,7 @@
 from sqlalchemy import create_engine, pool, func, make_url
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import sessionmaker
-from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Boolean, CHAR, JSON, UniqueConstraint, Float, Text
+from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Boolean, CHAR, JSON, UniqueConstraint, Float
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.sql.expression import true
 from contextlib import contextmanager
@@ -117,7 +117,6 @@ class ProcessedImage(Base):
     tries = Column(Integer, default=0)
 
 
-
 class StackFrame(Base):
     __tablename__ = 'stack_frames'
     id = Column(Integer, primary_key=True, autoincrement=True)

banzai/scheduling.py

@@ -1,5 +1,4 @@
 import os
-import traceback
 from datetime import datetime, timedelta, timezone
 from dateutil.parser import parse
 
@@ -303,4 +302,4 @@ def process_subframe(self, body: dict, runtime_context: dict):
     except Retry:
         raise
     except Exception:
-        logger.error(f'Error processing subframe: {traceback.format_exc()}')
+        logger.error("Error processing subframe: {error}".format(error=logs.format_exception()))

banzai/tests/site_e2e/conftest.py

@@ -6,6 +6,8 @@
 import subprocess
 import time
 from pathlib import Path
+from sqlalchemy import create_engine, text as sa_text
+from kombu import Connection, Queue
 
 import pytest
 
@@ -67,7 +69,6 @@ def poll_until(predicate, timeout, interval=5):
 
 def publish_to_queue(queue_name, body, broker_url='amqp://localhost:5672'):
     """Publish a JSON message to a RabbitMQ queue."""
-    from kombu import Connection, Queue
     with Connection(broker_url) as conn:
         queue = Queue(queue_name, channel=conn.channel())
         queue.declare()
@@ -142,7 +143,6 @@ def check():
 
 def wait_for_subscription_active(timeout=60):
     """Wait for the replication subscription to be enabled with an active worker."""
-    from sqlalchemy import create_engine, text as sa_text
 
     engine = create_engine(LOCAL_DB_ADDRESS)
 

banzai/tests/site_e2e/test_site_e2e.py

@@ -212,7 +212,6 @@ def test_07_reduction_completes(self, site_deployment):
     @pytest.mark.e2e_site_reduction
     def test_08_reduction_used_cached_calibrations(self, site_deployment):
         """Verify reduced frames used calibrations that exist in the local cache and DB."""
-        from astropy.io import fits
 
         output_dir = DATA_DIR / 'output'
         reduced_files = list(output_dir.rglob('*-e91.fits.fz'))
@@ -291,7 +290,6 @@ def test_11_cache_updated(self):
     @pytest.mark.e2e_site_reduction
     def test_12_subframe_stack_completes(self, site_deployment):
         """Verify subframe stacking processes a frame end-to-end."""
-        from astropy.io import fits
 
         raw_dir = DATA_DIR / 'raw'
         src_path = raw_dir / RAW_FRAME_FILENAME

banzai/tests/test_download_worker.py

@@ -128,7 +128,7 @@ def _attrs_for_type(cal_type, **overrides):
     return attrs
 
 
-def test_returns_top_2_per_config(db_address, tmp_path):
+def test_returns_top_2_per_config(db_address):
     inst_id = _seed_db(db_address)
     bias_attrs = _attrs_for_type('BIAS', configuration_mode='default', binning='1x1')
     with dbs.get_session(db_address) as session:
@@ -140,7 +140,7 @@ def test_returns_top_2_per_config(db_address, tmp_path):
     assert filenames == {'mid.fits', 'new.fits'}
 
 
-def test_partitions_independently_by_config(db_address, tmp_path):
+def test_partitions_independently_by_config(db_address):
     inst_id = _seed_db(db_address)
     with dbs.get_session(db_address) as session:
         for i, binning in enumerate(['1x1', '2x2']):
@@ -154,7 +154,7 @@ def test_partitions_independently_by_config(db_address, tmp_path):
                          'bias_2x2_1.fits', 'bias_2x2_2.fits'}
 
 
-def test_filters_by_instrument_type(db_address, tmp_path):
+def test_filters_by_instrument_type(db_address):
     sinistro_id = _seed_db(db_address, camera='fa01', inst_type='1m0-SciCam-Sinistro')
     floyds_id = _seed_db(db_address, camera='en01', inst_type='2m0-FLOYDS-SciCam')
     bias_attrs = _attrs_for_type('BIAS', configuration_mode='default', binning='1x1')
@@ -167,7 +167,7 @@ def test_filters_by_instrument_type(db_address, tmp_path):
     assert filenames == {'sinistro.fits'}
 
 
-def test_wildcard_returns_all_instrument_types(db_address, tmp_path):
+def test_wildcard_returns_all_instrument_types(db_address):
     sinistro_id = _seed_db(db_address, camera='fa01', inst_type='1m0-SciCam-Sinistro')
     floyds_id = _seed_db(db_address, camera='en01', inst_type='2m0-FLOYDS-SciCam')
     bias_attrs = _attrs_for_type('BIAS', configuration_mode='default', binning='1x1')
@@ -179,7 +179,7 @@ def test_wildcard_returns_all_instrument_types(db_address, tmp_path):
     assert filenames == {'sinistro.fits', 'floyds.fits'}
 
 
-def test_biases_ignore_filter(db_address, tmp_path):
+def test_biases_ignore_filter(db_address):
     """Biases taken with different filters should be grouped together."""
     inst_id = _seed_db(db_address)
     bias_attrs = _attrs_for_type('BIAS', configuration_mode='default', binning='1x1')
@@ -192,7 +192,7 @@ def test_biases_ignore_filter(db_address, tmp_path):
     assert filenames == {'bias_B.fits', 'bias_R.fits'}
 
 
-def test_darks_partitioned_by_temperature(db_address, tmp_path):
+def test_darks_partitioned_by_temperature(db_address):
     inst_id = _seed_db(db_address)
     with dbs.get_session(db_address) as session:
         for i, temp in enumerate(['5', '10']):
@@ -207,7 +207,7 @@ def test_darks_partitioned_by_temperature(db_address, tmp_path):
                          'dark_t10_1.fits', 'dark_t10_2.fits'}
 
 
-def test_skyflats_partitioned_by_filter(db_address, tmp_path):
+def test_skyflats_partitioned_by_filter(db_address):
     inst_id = _seed_db(db_address)
     with dbs.get_session(db_address) as session:
         for i, filt in enumerate(['V', 'B']):

banzai/tests/test_smart_stacking.py

@@ -13,6 +13,7 @@
                               push_notification, drain_notifications, REDIS_KEY_PREFIX,
                               process_notifications, finalize_stack, check_timeout,
                               discover_cameras, StackingSupervisor)
+from banzai.scheduling import process_subframe
 from banzai.main import SubframeListener
 
 pytestmark = pytest.mark.smart_stacking
@@ -48,7 +49,6 @@ def mock_redis():
 # Database operations
 # ---------------------------------------------------------------------------
 
-@pytest.mark.smart_stacking
 class TestDBOperations:
 
     def test_insert_and_query(self, db_address):
@@ -100,7 +100,6 @@ def test_update_stack_frame_filepath(self, db_address):
 # Status transitions
 # ---------------------------------------------------------------------------
 
-@pytest.mark.smart_stacking
 class TestStatusTransitions:
 
     def test_status_active_to_complete(self, db_address):
@@ -134,7 +133,6 @@ def test_status_active_to_timeout(self, db_address):
 # Timeout
 # ---------------------------------------------------------------------------
 
-@pytest.mark.smart_stacking
 class TestTimeout:
 
     def test_timeout_finalizes_stale_stacks(self, db_address):
@@ -154,7 +152,6 @@ def test_timeout_finalizes_stale_stacks(self, db_address):
 # Redis notifications
 # ---------------------------------------------------------------------------
 
-@pytest.mark.smart_stacking
 class TestRedisNotifications:
 
     def test_push_notification(self, mock_redis):
@@ -175,7 +172,6 @@ def test_drain_for_camera(self, mock_redis):
 # Multiple concurrent stacks
 # ---------------------------------------------------------------------------
 
-@pytest.mark.smart_stacking
 class TestConcurrentStacks:
 
     def test_concurrent_stacks_same_camera(self, db_address):
@@ -203,7 +199,6 @@ def test_concurrent_stacks_same_camera(self, db_address):
 # check_stack_complete
 # ---------------------------------------------------------------------------
 
-@pytest.mark.smart_stacking
 class TestCheckStackComplete:
 
     @staticmethod
@@ -237,7 +232,6 @@ def test_empty_frames(self):
 # Retention / cleanup
 # ---------------------------------------------------------------------------
 
-@pytest.mark.smart_stacking
 class TestRetention:
 
     def test_cleanup_old_records(self, db_address):
@@ -276,7 +270,6 @@ def test_cleanup_preserves_recent(self, db_address):
 # SubframeListener on_message
 # ---------------------------------------------------------------------------
 
-@pytest.mark.smart_stacking
 class TestSubframeListenerOnMessage:
     """on_message dispatches to Celery; no FITS I/O or DB work here."""
 
@@ -320,7 +313,6 @@ def test_on_message_invalid_no_dispatch(self, mock_task):
 # process_subframe Celery task
 # ---------------------------------------------------------------------------
 
-@pytest.mark.smart_stacking
 class TestProcessSubframe:
     """Test the Celery task that does the actual subframe processing."""
 
@@ -352,7 +344,6 @@ def _make_mock_image(output_dir='/data/processed', output_filename='frame-e09.fi
     ])
     @patch('banzai.scheduling.stage_utils.run_pipeline_stages')
     def test_process_subframe(self, mock_run_stages, last_frame_val, expected_is_last, db_address, mock_redis):
-        from banzai.scheduling import process_subframe
 
         mock_image = self._make_mock_image()
         mock_run_stages.return_value = [mock_image]
@@ -385,7 +376,6 @@ def test_process_subframe(self, mock_run_stages, last_frame_val, expected_is_las
 # Supervisor
 # ---------------------------------------------------------------------------
 
-@pytest.mark.smart_stacking
 class TestSupervisor:
 
     def test_discover_cameras(self, db_address):