ODS 完成

etl_billiards/tests/unit/task_test_utils.py

@@ -26,6 +26,7 @@ from tasks.refunds_task import RefundsTask
 from tasks.table_discount_task import TableDiscountTask
 from tasks.tables_task import TablesTask
 from tasks.topups_task import TopupsTask
+from utils.json_store import endpoint_to_filename
 
 DEFAULT_STORE_ID = 2790685415443269
 BASE_TS = "2025-01-01 10:00:00"
@@ -47,12 +48,6 @@ class TaskSpec:
         return endpoint_to_filename(self.endpoint)
 
 
-def endpoint_to_filename(endpoint: str) -> str:
-    """根据 API endpoint 生成稳定可复用的文件名，便于离线模式在目录中直接定位归档 JSON。"""
-    normalized = endpoint.strip("/").replace("/", "__").replace(" ", "_").lower()
-    return f"{normalized or 'root'}.json"
-
-
 def wrap_records(records: List[Dict], data_path: Sequence[str]):
     """按照 data_path 逐层包裹记录列表，使其结构与真实 API 返回体一致，方便离线回放。"""
     payload = records
@@ -140,6 +135,8 @@ class FakeDBOperations:
         self.commits = 0
         self.rollbacks = 0
         self.conn = FakeConnection()
+        # Pre-seeded query results (FIFO) to let tests control DB-returned rows
+        self.query_results: List[List[Dict]] = []
 
     def batch_upsert_with_returning(self, sql: str, rows: List[Dict], page_size: int = 1000):
         self.upserts.append(
@@ -167,6 +164,8 @@ class FakeDBOperations:
 
     def query(self, sql: str, params=None):
         self.executes.append({"sql": sql.strip(), "params": params, "type": "query"})
+        if self.query_results:
+            return self.query_results.pop(0)
         return []
 
     def cursor(self):