From d8b236047ebed446601958b65d860ed634e85c18 Mon Sep 17 00:00:00 2001
From: Lyn Nagara <lyn.nagara@gmail.com>
Date: Tue, 20 Jun 2023 11:15:18 -0700
Subject: [PATCH] test: Fix some more untyped tests (#4373)

---
 Makefile                   |  2 +-
 tests/state/test_cache.py  | 13 ++++++++-----
 tests/state/test_record.py |  2 +-
 3 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/Makefile b/Makefile
index 58312c876a..f7dcf92887 100644
--- a/Makefile
+++ b/Makefile
@@ -36,7 +36,7 @@ api-tests:
 	SNUBA_SETTINGS=test pytest -vv tests/*_api.py
 
 backend-typing:
-	mypy snuba tests scripts --strict --config-file mypy.ini --exclude 'tests/datasets|tests/query|tests/state|tests/snapshots|tests/clickhouse|tests/test_split.py'
+	mypy snuba tests scripts --strict --config-file mypy.ini --exclude 'tests/datasets|tests/query|tests/snapshots|tests/clickhouse|tests/test_split.py'
 
 install-python-dependencies:
 	pip uninstall -qqy uwsgi  # pip doesn't do well with swapping drop-ins
diff --git a/tests/state/test_cache.py b/tests/state/test_cache.py
index 295dfe93fa..fa4e90e8f8 100644
--- a/tests/state/test_cache.py
+++ b/tests/state/test_cache.py
@@ -16,7 +16,10 @@
 from snuba.state.cache.abstract import Cache, ExecutionError, ExecutionTimeoutError
 from snuba.state.cache.redis.backend import RedisCache
 from snuba.utils.codecs import ExceptionAwareCodec
-from snuba.utils.serializable_exception import SerializableException
+from snuba.utils.serializable_exception import (
+    SerializableException,
+    SerializableExceptionDict,
+)
 from tests.assertions import assert_changes, assert_does_not_change
 
 redis_client = get_redis_client(RedisClientKey.CACHE)
@@ -44,7 +47,7 @@ def encode(self, value: bytes) -> bytes:
 
     def decode(self, value: bytes) -> bytes:
         try:
-            ret = rapidjson.loads(value)
+            ret: SerializableExceptionDict = rapidjson.loads(value)
             if not isinstance(ret, dict):
                 return value
             if ret.get("__type__", "NOP") == "SerializableException":
@@ -55,7 +58,7 @@ def decode(self, value: bytes) -> bytes:
         return value
 
     def encode_exception(self, value: SerializableException) -> bytes:
-        return rapidjson.dumps(value.to_dict()).encode("utf-8")
+        return bytes(rapidjson.dumps(value.to_dict()).encode("utf-8"))
 
 
 @pytest.fixture
@@ -260,10 +263,10 @@ def test_notify_queue_ttl() -> None:
     num_waiters = 9
 
     class DelayedRedisClient:
-        def __init__(self, redis_client):
+        def __init__(self, redis_client: RedisClientType) -> None:
             self._client = redis_client
 
-        def __getattr__(self, attr: str):
+        def __getattr__(self, attr: str) -> Any:
             # simulate the queue pop taking longer than expected.
             # the notification queue TTL is 60 seconds so running into a timeout
             # shouldn't happen (unless something has drastically changed in the TTL
diff --git a/tests/state/test_record.py b/tests/state/test_record.py
index f9cad113eb..f22e4e8470 100644
--- a/tests/state/test_record.py
+++ b/tests/state/test_record.py
@@ -1,5 +1,5 @@
 from snuba.state import _kafka_producer
 
 
-def test_get_producer():
+def test_get_producer() -> None:
     assert _kafka_producer() is not None