Improvements to tests

oonipipeline/tests/test_temporal_e2e.py

@@ -13,6 +13,7 @@
 )
 from oonipipeline.temporal.workers import ACTIVTIES
 
+from .utils import wait_for_mutations
 
 @pytest.mark.asyncio
 async def test_observation_workflow(datadir, db):
@@ -71,8 +72,8 @@ async def test_observation_workflow(datadir, db):
                 id="obs-wf-2",
                 task_queue=TASK_QUEUE_NAME,
             )
-            db.execute("OPTIMIZE TABLE buffer_obs_web")
             db.execute("OPTIMIZE TABLE obs_web")
+            wait_for_mutations(db, "obs_web")
             res = db.execute(
                 """
                 SELECT bucket_date,

oonipipeline/tests/test_workflows.py

@@ -4,7 +4,6 @@
 import sqlite3
 from typing import List, Tuple
 from unittest.mock import MagicMock
-import time
 
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import Worker
@@ -21,18 +20,14 @@
 from oonipipeline.temporal.activities.common import (
     ClickhouseParams,
     OptimizeTablesParams,
-    UpdateAssetsParams,
     get_obs_count_by_cc,
     ObsCountParams,
 )
 from oonipipeline.temporal.activities.observations import (
     DeletePreviousRangeParams,
     GetPreviousRangeParams,
-    MakeObservationsFileEntryBatch,
     MakeObservationsParams,
     MakeObservationsResult,
-    ObservationBatches,
-    make_observations_for_file_entry,
     make_observations_for_file_entry_batch,
 )
 from oonipipeline.transforms.measurement_transformer import MeasurementTransformer
@@ -63,19 +58,7 @@
 )
 from oonipipeline.temporal.workflows.common import TASK_QUEUE_NAME
 
-# from oonipipeline.workflows.response_archiver import ResponseArchiver
-# from oonipipeline.workflows.fingerprint_hunter import fingerprint_hunter
-
-
-def wait_for_mutations(db, table_name):
-    while True:
-        res = db.execute(
-            f"SELECT * FROM system.mutations WHERE is_done=0 AND table='{table_name}';"
-        )
-        if len(res) == 0:  # type: ignore
-            break
-        time.sleep(1)
-
+from .utils import wait_for_mutations
 
 def test_get_prev_range(db):
     db.execute("DROP TABLE IF EXISTS test_range")

oonipipeline/tests/utils.py

@@ -0,0 +1,15 @@
+# from oonipipeline.workflows.response_archiver import ResponseArchiver
+# from oonipipeline.workflows.fingerprint_hunter import fingerprint_hunter
+
+
+import time
+
+
+def wait_for_mutations(db, table_name):
+    while True:
+        res = db.execute(
+            f"SELECT * FROM system.mutations WHERE is_done=0 AND table='{table_name}';"
+        )
+        if len(res) == 0:  # type: ignore
+            break
+        time.sleep(1)