Merges pull request #3

Off of Amazon!!!

1 files changed, 14 insertions, 12 deletions

diff --git a/tests/http/test_notifications.py b/tests/http/test_notifications.py
index 508b236..3df2423 100644
--- a/tests/http/test_notifications.py
+++ b/tests/http/test_notifications.py
@@ -1,14 +1,15 @@
-import pytest
 import json
-import redis
-from typing import Dict, Any
-from httpx import AsyncClient
+from typing import Any
 from uuid import uuid4
 
+import pytest
+from generalresearch.redis_helper import RedisConfig
+from httpx import AsyncClient
+
 from jb.config import JB_EVENTS_STREAM, settings
+from jb.models.assignment import AssignmentStub
 from jb.models.event import MTurkEvent
 from jb.models.hit import Hit
-from jb.models.assignment import AssignmentStub
 
 
 class TestNotifications:
@@ -51,13 +52,14 @@ class TestNotifications:
     @pytest.mark.anyio
     async def test_mturk_notifications(
         self,
-        redis: redis.Redis,
+        redis_config: RedisConfig,
         httpxclient: AsyncClient,
         hit_record: Hit,
         assignment_stub_record: AssignmentStub,
-        mturk_event_body_record: Dict[str, Any],
+        mturk_event_body_record: dict[str, Any],
     ):
         client = httpxclient
+        redis_client = redis_config.create_redis_client()
 
         json_msg = json.loads(mturk_event_body_record["Message"])
         # Assert the mturk event is owned by the correct account
@@ -76,7 +78,7 @@ class TestNotifications:
         )
 
         # Confirm the stream is empty
-        assert redis.xlen(JB_EVENTS_STREAM) == 0
+        assert redis_client.xlen(JB_EVENTS_STREAM) == 0
 
         res = await client.post(
             url=f"/{settings.sns_path}/", json=mturk_event_body_record
@@ -85,20 +87,20 @@ class TestNotifications:
 
         # Now that we POSTed, confirm the stream has 1 event in it
         # Confirm the stream is empty
-        assert redis.xlen(JB_EVENTS_STREAM) == 1
+        assert redis_client.xlen(JB_EVENTS_STREAM) == 1
 
         # AMT SNS needs to receive a 200 response to stop retrying the notification
         assert res.status_code == 200
         assert res.json() == {"status": "ok"}
 
         # Check that the event was enqueued in Redis
-        msg_res = redis.xread(streams={JB_EVENTS_STREAM: 0}, count=1, block=100)
+        msg_res = redis_client.xread(streams={JB_EVENTS_STREAM: 0}, count=1, block=100)
         msg_res = msg_res[0][1][0]
         msg_id, msg = msg_res
-        redis.xdel(JB_EVENTS_STREAM, msg_id)
+        redis_client.xdel(JB_EVENTS_STREAM, msg_id)
 
         # After running xdel, we can confirm the stream is empty
-        assert redis.xlen(JB_EVENTS_STREAM) == 0
+        assert redis_client.xlen(JB_EVENTS_STREAM) == 0
 
         msg_json = msg["data"]
         event = MTurkEvent.model_validate_json(msg_json)