1 files changed, 57 insertions, 19 deletions
diff --git a/tests/managers/thl/test_user_manager/test_base.py b/tests/managers/thl/test_user_manager/test_base.py
index 0d7ffef..5d12052 100644
--- a/tests/managers/thl/test_user_manager/test_base.py
+++ b/tests/managers/thl/test_user_manager/test_base.py
@@ -1,23 +1,37 @@
 import logging
-from datetime import datetime, timezone
+from collections.abc import Callable
+from datetime import UTC, datetime
 from random import randint
+from typing import TYPE_CHECKING
 from uuid import uuid4
 
 import pytest
 
 from generalresearch.managers.thl.user_manager import (
-    UserCreateNotAllowedError,
     get_bp_user_create_limit_hourly,
 )
+from generalresearch.managers.thl.user_manager.exceptions import (
+    UserCreateNotAllowedError,
+)
+from generalresearch.managers.thl.user_manager.mysql_user_manager import (
+    MysqlUserManager,
+)
 from generalresearch.managers.thl.user_manager.rate_limit import (
     RateLimitItemPerHourConstantKey,
+    UserManagerLimiter,
 )
-from generalresearch.managers.thl.user_manager.user_manager import (
-    UserManager,
-)
-from generalresearch.models.thl.product import Product, UserCreateConfig
+from generalresearch.models.thl.product import UserCreateConfig
 from generalresearch.models.thl.user import User
 
+if TYPE_CHECKING:
+    from generalresearch.managers.thl.product import ProductManager
+    from generalresearch.managers.thl.user_manager.user_manager import (
+        UserManager,
+    )
+    from generalresearch.managers.thl.userhealth import AuditLogManager
+    from generalresearch.models.thl.product import Product
+    from generalresearch.pg_helper import PostgresConfig
+
 logger = logging.getLogger()
 
 
@@ -83,10 +97,11 @@ class TestUserManager:
 
 class TestBlockUserManager:
 
-    def test_block_user(self, product, user_manager: UserManager):
+    def test_block_user(self, product: Product, user_manager: UserManager):
         product_user_id = f"user-{uuid4().hex[:10]}"
 
         # mysql_user_manager to skip user creation limit check
+        assert isinstance(user_manager.mysql_user_manager, MysqlUserManager)
         user: User = user_manager.mysql_user_manager.create_user(
             product_id=product.id, product_user_id=product_user_id
         )
@@ -109,16 +124,19 @@ class TestBlockUserManager:
         user = user_manager.get_user(user_id=user.user_id)
         assert user.blocked
 
-    def test_block_user_whitelist(self, product, user_manager, thl_web_rw):
+    def test_block_user_whitelist(
+        self, product: Product, user_manager: UserManager, thl_web_rw: PostgresConfig
+    ):
         product_user_id = f"user-{uuid4().hex[:10]}"
 
         # mysql_user_manager to skip user creation limit check
+        assert isinstance(user_manager.mysql_user_manager, MysqlUserManager)
         user: User = user_manager.mysql_user_manager.create_user(
             product_id=product.id, product_user_id=product_user_id
         )
         assert not user.blocked
 
-        now = datetime.now(tz=timezone.utc)
+        now = datetime.now(tz=UTC)
         # Adds user to whitelist
         thl_web_rw.execute_write(
             """
@@ -135,8 +153,13 @@ class TestBlockUserManager:
 
 class TestCreateUserManager:
 
-    def test_create_user(self, product_manager, thl_web_rw, user_manager):
-        product: Product = product_manager.create_dummy(
+    def test_create_user(
+        self,
+        product_factory: Callable[..., Product],
+        thl_web_rw: PostgresConfig,
+        user_manager: UserManager,
+    ):
+        product: Product = product_factory(
             user_create_config=UserCreateConfig(
                 min_hourly_create_limit=10, max_hourly_create_limit=69
             ),
@@ -144,6 +167,7 @@ class TestCreateUserManager:
 
         product_user_id = f"user-{uuid4().hex[:10]}"
 
+        assert isinstance(user_manager.mysql_user_manager, MysqlUserManager)
         user: User = user_manager.mysql_user_manager.create_user(
             product_id=product.id, product_user_id=product_user_id
         )
@@ -156,7 +180,7 @@ class TestCreateUserManager:
 
         # make sure thl_user row is created
         res_thl_user = thl_web_rw.execute_sql_query(
-            query=f"""
+            query="""
                 SELECT * 
                 FROM thl_user AS u
                 WHERE u.id = %s
@@ -172,8 +196,13 @@ class TestCreateUserManager:
         assert u2.user_id == user.user_id
         assert u2.uuid == user.uuid
 
-    def test_create_user_integrity_error(self, product_manager, user_manager, caplog):
-        product: Product = product_manager.create_dummy(
+    def test_create_user_integrity_error(
+        self,
+        user_manager: UserManager,
+        product_factory: Callable[..., Product],
+        caplog,
+    ):
+        product: Product = product_factory(
             product_id=uuid4().hex,
             team_id=uuid4().hex,
             name=f"Test Product ID #{uuid4().hex[:6]}",
@@ -185,6 +214,7 @@ class TestCreateUserManager:
         product_user_id = f"user-{uuid4().hex[:10]}"
         rand_msg = f"log-{uuid4().hex}"
 
+        assert isinstance(user_manager.mysql_user_manager, MysqlUserManager)
         with caplog.at_level(logging.INFO):
             logger.info(rand_msg)
             user1 = user_manager.mysql_user_manager.create_user(
@@ -213,9 +243,14 @@ class TestCreateUserManager:
 
         assert user1 == user2
 
-    def test_raise_allow_user_create(self, product_manager, user_manager):
+    def test_raise_allow_user_create(
+        self,
+        product_manager: ProductManager,
+        user_manager: UserManager,
+        product_factory: Callable[..., Product],
+    ):
         rand_num = randint(25, 200)
-        product: Product = product_manager.create_dummy(
+        product: Product = product_factory(
             product_id=uuid4().hex,
             team_id=uuid4().hex,
             name=f"Test Product ID #{uuid4().hex[:6]}",
@@ -247,10 +282,11 @@ class TestCreateUserManager:
         assert key == f"LIMITER/thl-grpc/allow_user_create/{instance.id}"
 
         # make sure we clear the key or subsequent tests will fail
+        assert isinstance(user_manager.user_manager_limiter, UserManagerLimiter)
         user_manager.user_manager_limiter.storage.clear(key=key)
 
         n = 0
-        with pytest.raises(expected_exception=UserCreateNotAllowedError) as cm:
+        with pytest.raises(expected_exception=UserCreateNotAllowedError):
             for n, _ in enumerate(range(rl_value + 5)):
                 user_manager.user_manager_limiter.raise_allow_user_create(
                     product=product
@@ -260,14 +296,16 @@ class TestCreateUserManager:
 
 class TestUserManagerMethods:
 
-    def test_audit_log(self, user_manager, user, audit_log_manager):
+    def test_audit_log(
+        self, user_manager: UserManager, user: User, audit_log_manager: AuditLogManager
+    ):
         from generalresearch.models.thl.userhealth import AuditLog
 
         res = audit_log_manager.filter_by_user_id(user_id=user.user_id)
         assert len(res) == 0
 
         msg = uuid4().hex
-        user_manager.audit_log(user=user, level=30, event_type=msg)
+        user_manager.audit_log(audit_log_manager, user=user, level=30, event_type=msg)
 
         res = audit_log_manager.filter_by_user_id(user_id=user.user_id)
         assert len(res) == 1