1 files changed, 4 insertions, 2 deletions
diff --git a/tests/replication/test_account_validity.py b/tests/replication/test_account_validity.py
index 408eb266b9..7353bbce96 100644
--- a/tests/replication/test_account_validity.py
+++ b/tests/replication/test_account_validity.py
@@ -33,6 +33,8 @@ class MockAccountValidityStore:
     ):
         self._api = api
 
+        api.register_cached_function(self.is_user_expired)
+
     async def create_db(self):
         def create_table_txn(txn: LoggingTransaction):
             txn.execute(
@@ -88,13 +90,13 @@ class MockAccountValidityStore:
                 ),
             )
 
-            txn.call_after(self.is_user_expired.invalidate, (user_id,))
-
         await self._api.run_db_interaction(
             "account_validity_set_expired_user",
             set_expired_user_txn,
         )
 
+        await self._api.invalidate_cache(self.is_user_expired, (user_id,))
+
 
 class MockAccountValidity:
     def __init__(