1 files changed, 45 insertions, 12 deletions
diff --git a/tests/managers/thl/test_ipinfo.py b/tests/managers/thl/test_ipinfo.py
index c89312b..48b9efd 100644
--- a/tests/managers/thl/test_ipinfo.py
+++ b/tests/managers/thl/test_ipinfo.py
@@ -1,3 +1,5 @@
+from collections.abc import Callable
+
 import faker
 
 from generalresearch.managers.thl.ipinfo import (
@@ -5,14 +7,22 @@ from generalresearch.managers.thl.ipinfo import (
     IPGeonameManager,
     IPInformationManager,
 )
-from generalresearch.models.thl.ipinfo import IPGeoname, IPInformation
+from generalresearch.models.thl.ipinfo import (
+    GeoIPInformation,
+    IPGeoname,
+    IPInformation,
+)
+from generalresearch.pg_helper import PostgresConfig
+from generalresearch.redis_helper import RedisConfig
 
 fake = faker.Faker()
 
 
 class TestIPGeonameManager:
 
-    def test_init(self, thl_web_rr: PostgresConfig, ip_geoname_manager: IPGeonameManager):
+    def test_init(
+        self, thl_web_rr: PostgresConfig, ip_geoname_manager: IPGeonameManager
+    ):
 
         instance = IPGeonameManager(pg_config=thl_web_rr)
         assert isinstance(instance, IPGeonameManager)
@@ -31,7 +41,9 @@ class TestIPGeonameManager:
 
 class TestIPInformationManager:
 
-    def test_init(self, thl_web_rr: PostgresConfig, ip_information_manager: IPInformationManager):
+    def test_init(
+        self, thl_web_rr: PostgresConfig, ip_information_manager: IPInformationManager
+    ):
         instance = IPInformationManager(pg_config=thl_web_rr)
         assert isinstance(instance, IPInformationManager)
         assert isinstance(ip_information_manager, IPInformationManager)
@@ -45,7 +57,12 @@ class TestIPInformationManager:
 
         assert res[0].model_dump_json() == instance.model_dump_json()
 
-    def test_prefetch_geoname(self, ip_information, ip_geoname, thl_web_rr):
+    def test_prefetch_geoname(
+        self,
+        ip_information: IPInformation,
+        ip_geoname: IPGeoname,
+        thl_web_rr: PostgresConfig,
+    ):
         assert isinstance(ip_information, IPInformation)
 
         assert ip_information.geoname_id == ip_geoname.geoname_id
@@ -62,11 +79,16 @@ class TestGeoIpInfoManager:
         thl_redis_config: RedisConfig,
         geoipinfo_manager: GeoIpInfoManager,
     ):
-        instance = GeoIpInfoManager(pg_config=thl_web_rr: PostgresConfig, redis_config=thl_redis_config)
+        instance = GeoIpInfoManager(pg_config=thl_web_rr, redis_config=thl_redis_config)
         assert isinstance(instance, GeoIpInfoManager)
         assert isinstance(geoipinfo_manager, GeoIpInfoManager)
 
-    def test_multi(self, ip_information_factory, ip_geoname, geoipinfo_manager):
+    def test_multi(
+        self,
+        ip_information_factory: Callable[..., IPInformation],
+        ip_geoname: IPGeoname,
+        geoipinfo_manager: GeoIpInfoManager,
+    ):
         ip = fake.ipv4_public()
         ip_information_factory(ip=ip, geoname=ip_geoname)
         ips = [ip]
@@ -93,7 +115,12 @@ class TestGeoIpInfoManager:
         assert res[ip] is not None
         assert res[ip2] is not None
 
-    def test_multi_ipv6(self, ip_information_factory, ip_geoname, geoipinfo_manager):
+    def test_multi_ipv6(
+        self,
+        ip_information_factory: Callable[..., IPInformation],
+        ip_geoname: IPGeoname,
+        geoipinfo_manager: GeoIpInfoManager,
+    ):
         ip = fake.ipv6()
         # Make another IP that will be in the same /64 block.
         ip2 = ip[:-1] + "a" if ip[-1] != "a" else ip[:-1] + "b"
@@ -108,13 +135,19 @@ class TestGeoIpInfoManager:
         # Looks up in redis, if not exists, looks in mysql, then sets
         #   the caches that didn't exist.
         res = geoipinfo_manager.get_multi(ip_addresses=ips)
-        assert res[ip].ip == ip
-        assert res[ip].lookup_prefix == "/64"
-        assert res[ip2].ip == ip2
-        assert res[ip2].lookup_prefix == "/64"
+
+        res1 = res[ip]
+        assert isinstance(res1, GeoIPInformation)
+        assert res1.ip == ip
+        assert res1.lookup_prefix == "/64"
+
+        res2 = res[ip2]
+        assert isinstance(res2, GeoIPInformation)
+        assert res2.ip == ip2
+        assert res2.lookup_prefix == "/64"
         # they should be the same basically, except for the ip
 
-    def test_doesnt_exist(self, geoipinfo_manager):
+    def test_doesnt_exist(self, geoipinfo_manager: GeoIpInfoManager):
         ip = fake.ipv4_public()
         res = geoipinfo_manager.get_multi(ip_addresses=[ip])
         assert res == {ip: None}