diff --git a/stubs/asynctest/__init__.pyi b/stubs/asynctest/__init__.pyi
new file mode 100644
index 0000000..e69de29
diff --git a/stubs/asynctest/mock/__init__.pyi b/stubs/asynctest/mock/__init__.pyi
new file mode 100644
index 0000000..b1047a7
--- /dev/null
+++ b/stubs/asynctest/mock/__init__.pyi
@@ -0,0 +1,2 @@
+from unittest.mock import *
+from unittest.mock import Mock as CoroutineMock
diff --git a/tests/async_mock.py b/tests/async_mock.py
deleted file mode 100644
index 6e897bc..0000000
--- a/tests/async_mock.py
+++ /dev/null
@@ -1,7 +0,0 @@
-import sys
-
-if sys.version_info[:2] < (3, 8):
-    from asynctest.mock import *  # noqa
-    from asynctest.mock import CoroutineMock as AsyncMock  # noqa
-else:
-    from unittest.mock import *  # noqa
diff --git a/tests/async_mock/__init__.py b/tests/async_mock/__init__.py
new file mode 100644
index 0000000..7ed6459
--- /dev/null
+++ b/tests/async_mock/__init__.py
@@ -0,0 +1,14 @@
+import sys
+
+if sys.version_info[:2] < (3, 8):
+    from asynctest.mock import (
+        Mock,
+        NonCallableMock,
+        CoroutineMock as AsyncMock,
+        call,
+        patch,
+    )
+else:
+    from unittest.mock import Mock, NonCallableMock, call, patch
+
+__all__ = ["Mock", "AsyncMock", "NonCallableMock", "call", "patch"]
diff --git a/tests/types.py b/tests/helper_types.py
similarity index 100%
rename from tests/types.py
rename to tests/helper_types.py
diff --git a/tests/test_entities.py b/tests/test_entities.py
index 15f9400..a70a258 100644
--- a/tests/test_entities.py
+++ b/tests/test_entities.py
@@ -8,12 +8,11 @@
 from homeassistant.core import HomeAssistant
 
 from tests.async_mock import Mock, NonCallableMock, call, patch
-from tests.types import GetHub
 
 
 @patch("custom_components.hubitat.entities.get_hub")
 @patch("custom_components.hubitat.entities.entity_registry")
-async def test_entity_migration(get_hub: GetHub, entity_registry) -> None:
+async def test_entity_migration(get_hub: Mock, entity_registry: Mock) -> None:
     mock_device_1 = NonCallableMock(type="switch", attributes=["state"])
     mock_device_2 = NonCallableMock(type="fan", attributes=["state"])
     MockHub = Mock(spec=Hub)
@@ -63,7 +62,7 @@ def _is_switch(device: Device) -> bool:
 
 @patch("custom_components.hubitat.entities.get_hub")
 @patch("custom_components.hubitat.entities.HubitatEventEmitter")
-async def test_add_event_emitters(HubitatEventEmitter, get_hub) -> None:
+async def test_add_event_emitters(HubitatEventEmitter: Mock, get_hub: Mock) -> None:
     mock_device_1 = NonCallableMock(type="switch", attributes=["state"])
     mock_device_2 = NonCallableMock(type="button", attributes=["state"])
     MockHub = Mock(spec=Hub)
diff --git a/tests/test_lock.py b/tests/test_lock.py
index 1deed18..e7104cc 100644
--- a/tests/test_lock.py
+++ b/tests/test_lock.py
@@ -1,20 +1,24 @@
 from hubitatmaker.const import ATTR_LOCK_CODES
 from hubitatmaker.types import Attribute
 
-from tests.async_mock import MagicMock
+from tests.async_mock import Mock
 
 
 def test_normal_lock_codes() -> None:
-    hub = MagicMock()
-    device = MagicMock()
-    device.attributes = {
-        ATTR_LOCK_CODES: Attribute(
-            {
-                "name": ATTR_LOCK_CODES,
-                "currentValue": '{"1":{"name":"Test","code":"1234"}}',
-            }
-        )
-    }
+    hub = Mock()
+    hub.configure_mock(token="abc1235Qbxyz")
+
+    device = Mock()
+    device.configure_mock(
+        attributes={
+            ATTR_LOCK_CODES: Attribute(
+                {
+                    "name": ATTR_LOCK_CODES,
+                    "currentValue": '{"1":{"name":"Test","code":"1234"}}',
+                }
+            )
+        }
+    )
 
     from custom_components.hubitat.lock import HubitatLock
 
@@ -23,13 +27,17 @@ def test_normal_lock_codes() -> None:
 
 
 def test_encrypted_lock_codes() -> None:
-    hub = MagicMock()
-    device = MagicMock()
-    device.attributes = {
-        ATTR_LOCK_CODES: Attribute(
-            {"name": ATTR_LOCK_CODES, "currentValue": "abc1235Qbxyz"}
-        )
-    }
+    hub = Mock()
+    hub.configure_mock(token="abc1235Qbxyz")
+
+    device = Mock()
+    device.configure_mock(
+        attributes={
+            ATTR_LOCK_CODES: Attribute(
+                {"name": ATTR_LOCK_CODES, "currentValue": "abc1235Qbxyz"}
+            )
+        }
+    )
 
     from custom_components.hubitat.lock import HubitatLock