1 files changed, 17 insertions, 6 deletions
diff --git a/tests/middleware/test_authenticate.py b/tests/middleware/test_authenticate.py
index 2acfe92..9fccb32 100644
--- a/tests/middleware/test_authenticate.py
+++ b/tests/middleware/test_authenticate.py
@@ -34,7 +34,7 @@ def test_unauthenticated_request(app, method):
 def test_unauthenticated_request_with_key(app, method):
     environ = {
         "REQUEST_METHOD": method,
-        "paste.store": None,
+        "paste.db_conn": None,
         "HTTP_AUTHORIZATION": "ApiKey AAAA",
     }
     response = call_app(app, environ)
@@ -64,13 +64,22 @@ def test_authenticate_malformed_key(app, method, key):
     assert ("WWW-Authenticate", "APIKey") in response.headers
 
 
-class MockDB:
+class MockConnection:
     def __init__(self, check_token):
         self.check_token = check_token
 
 
+class MockStore:
+    def __init__(self, c):
+        assert isinstance(c, MockConnection)
+        self.conn = c
+
+    def check_token(self, tok):
+        return self.conn.check_token(tok)
+
+
 @pytest.mark.parametrize("method", ["POST", "PUT", "DELETE"])
-def test_authenticate_check_token_fail(app, method):
+def test_authenticate_check_token_fail(app, method, monkeypatch):
     check_token_called = False
     token = b"test"
 
@@ -82,10 +91,11 @@ def test_authenticate_check_token_fail(app, method):
 
     environ = {
         "REQUEST_METHOD": method,
-        "paste.store": MockDB(check_token),
+        "paste.db_conn": MockConnection(check_token),
         "HTTP_AUTHORIZATION": f"APIKey {b64encode(token).decode()}",
     }
 
+    monkeypatch.setattr("paste.Store", MockStore)
     response = call_app(app, environ)
     assert check_token_called
     assert response.data == b"401 Unauthorized\n"
@@ -95,7 +105,7 @@ def test_authenticate_check_token_fail(app, method):
 
 
 @pytest.mark.parametrize("method", ["POST", "PUT", "DELETE"])
-def test_authenticate_check_token_success(app, method):
+def test_authenticate_check_token_success(app, method, monkeypatch):
     check_token_called = False
     token = b"test"
 
@@ -107,10 +117,11 @@ def test_authenticate_check_token_success(app, method):
 
     environ = {
         "REQUEST_METHOD": method,
-        "paste.store": MockDB(check_token),
+        "paste.db_conn": MockConnection(check_token),
         "HTTP_AUTHORIZATION": f"APIKey {b64encode(token).decode()}",
     }
 
+    monkeypatch.setattr("paste.Store", MockStore)
     response = call_app(app, environ)
     assert check_token_called
     assert response.data == b"Hello, world!"