1 files changed, 41 insertions, 43 deletions

diff --git a/tests/incite/mergers/test_ym_survey_merge.py b/tests/incite/mergers/test_ym_survey_merge.py
index 4c2df6b..d83a98c 100644
--- a/tests/incite/mergers/test_ym_survey_merge.py
+++ b/tests/incite/mergers/test_ym_survey_merge.py
@@ -1,25 +1,28 @@
-from datetime import timedelta, timezone, datetime
+from __future__ import annotations
+
+from collections.abc import Callable
+from datetime import UTC, datetime, timedelta
 from itertools import product
+from typing import TYPE_CHECKING
 
 import pandas as pd
 import pytest
+from dask.distributed import Client as DaskClient
+
+if TYPE_CHECKING:
+    from generalresearch.incite.collections.thl_web import (
+        SessionDFCollection,
+        WallDFCollection,
+    )
+    from generalresearch.incite.mergers.foundations.enriched_session import (
+        EnrichedSessionMerge,
+    )
+    from generalresearch.incite.mergers.ym_survey_wall import YMSurveyWallMerge
+    from generalresearch.models.thl.product import Product
+    from generalresearch.models.thl.user import User
+    from generalresearch.pg_helper import PostgresConfig
 
 # noinspection PyUnresolvedReferences
-from distributed.utils_test import (
-    gen_cluster,
-    client_no_amm,
-    loop,
-    loop_in_thread,
-    cleanup,
-    cluster_fixture,
-    client,
-)
-
-from test_utils.incite.collections.conftest import wall_collection, session_collection
-from test_utils.incite.mergers.conftest import (
-    enriched_session_merge,
-    ym_survey_wall_merge,
-)
 
 
 @pytest.mark.parametrize(
@@ -28,11 +31,7 @@ from test_utils.incite.mergers.conftest import (
         product(
             ["12h", "3D"],
             [timedelta(days=30)],
-            [
-                (datetime.now(tz=timezone.utc) - timedelta(days=35)).replace(
-                    microsecond=0
-                )
-            ],
+            [(datetime.now(tz=UTC) - timedelta(days=35)).replace(microsecond=0)],
         )
     ),
 )
@@ -46,18 +45,17 @@ class TestYMSurveyMerge:
 
     def test_base(
         self,
-        client_no_amm,
-        user_factory,
-        product,
-        ym_survey_wall_merge,
-        wall_collection,
-        session_collection,
-        enriched_session_merge,
-        delete_df_collection,
-        incite_item_factory,
-        thl_web_rr,
+        client_no_amm: DaskClient,
+        user_factory: Callable[..., User],
+        product: Product,
+        ym_survey_wall_merge: YMSurveyWallMerge,
+        wall_collection: WallDFCollection,
+        session_collection: SessionDFCollection,
+        enriched_session_merge: EnrichedSessionMerge,
+        delete_df_collection: Callable[..., None],
+        incite_item_factory: Callable[..., None],
+        thl_web_rr: PostgresConfig,
     ):
-        from generalresearch.models.thl.user import User
 
         delete_df_collection(coll=session_collection)
         user: User = user_factory(product=product, created=session_collection.start)
@@ -85,10 +83,10 @@ class TestYMSurveyMerge:
         assert enriched_session_merge.progress.has_archive.eq(True).all()
 
         ddf = enriched_session_merge.ddf()
-        df: pd.DataFrame = client_no_amm.compute(collections=ddf, sync=True)
+        df1: pd.DataFrame | None = client_no_amm.compute(collections=ddf, sync=True)
 
-        assert isinstance(df, pd.DataFrame)
-        assert not df.empty
+        assert isinstance(df1, pd.DataFrame)
+        assert not df1.empty
 
         # --
 
@@ -102,18 +100,18 @@ class TestYMSurveyMerge:
         # --
 
         ddf = ym_survey_wall_merge.ddf()
-        df: pd.DataFrame = client_no_amm.compute(collections=ddf, sync=True)
+        df2: pd.DataFrame | None = client_no_amm.compute(collections=ddf, sync=True)
 
-        assert isinstance(df, pd.DataFrame)
-        assert not df.empty
+        assert isinstance(df2, pd.DataFrame)
+        assert not df2.empty
 
         # --
-        assert df.product_id.nunique() == 1
-        assert df.team_id.nunique() == 1
-        assert df.source.nunique() > 1
+        assert df2.product_id.nunique() == 1
+        assert df2.team_id.nunique() == 1
+        assert df2.source.nunique() > 1
 
-        started_min_ts = df.started.min()
-        started_max_ts = df.started.max()
+        started_min_ts = df2.started.min()
+        started_max_ts = df2.started.max()
 
         assert type(started_min_ts) is pd.Timestamp
         assert type(started_max_ts) is pd.Timestamp