fix tests

Author: Swatinem

database/base.py

@@ -2,8 +2,7 @@
 
 from sqlalchemy import Column, types
 from sqlalchemy.dialects.postgresql import UUID
-from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import Session, declarative_base
 
 from helpers.clock import get_utc_now
 

database/engine.py

@@ -54,7 +54,7 @@ def create_session(self):
             timeseries_engine = self.timeseries_engine
 
             class RoutingSession(Session):
-                def get_bind(self, mapper=None, clause=None):
+                def get_bind(self, mapper=None, clause=None, **kwargs):
                     if mapper is not None and issubclass(
                         mapper.class_, TimeseriesBaseModel
                     ):

ta_storage/pg.py

@@ -37,14 +37,15 @@ def get_repo_flag_ids(db_session: Session, repoid: int, flags: list[str]) -> set
     if not flags:
         return set()
 
-    return set(
-        db_session.query(RepositoryFlag.id_)
+    return {
+        flag.id_
+        for flag in db_session.query(RepositoryFlag.id_)
         .filter(
             RepositoryFlag.repository_id == repoid,
             RepositoryFlag.flag_name.in_(flags),
         )
         .all()
-    )
+    }
 
 
 def modify_structures(