Rejig testcase to make it more extensible

1 files changed, 22 insertions, 17 deletions

diff --git a/tests/rest/client/v1/test_admin.py b/tests/rest/client/v1/test_admin.py
index 2b7ade9f62..fb4ac6b95f 100644
--- a/tests/rest/client/v1/test_admin.py
+++ b/tests/rest/client/v1/test_admin.py
@@ -366,38 +366,43 @@ class ShutdownRoomTestCase(unittest.HomeserverTestCase):
         self.event_creation_handler = hs.get_event_creation_handler()
         hs.config.user_consent_version = "1"
 
-        self._consent_uri_builder = Mock()
-        self._consent_uri_builder.build_user_consent_uri.return_value = (
+        consent_uri_builder = Mock()
+        consent_uri_builder.build_user_consent_uri.return_value = (
             "http://example.com"
         )
+        self.event_creation_handler._consent_uri_builder = consent_uri_builder
 
         self.store = hs.get_datastore()
 
-    def test_shutdown_room_conset(self):
-        admin_user = self.register_user("admin", "pass", admin=True)
-        admin_user_tok = self.login("admin", "pass")
+        self.admin_user = self.register_user("admin", "pass", admin=True)
+        self.admin_user_tok = self.login("admin", "pass")
+
+        self.other_user = self.register_user("user", "pass")
+        self.other_user_token = self.login("user", "pass")
+
+        # Mark the admin user as having consented
+        self.get_success(
+            self.store.user_set_consent_version(self.admin_user, "1"),
+        )
 
-        other_user = self.register_user("user", "pass")
-        other_user_token = self.login("user", "pass")
+    def test_shutdown_room_conset(self):
+        self.event_creation_handler._block_events_without_consent_error = None
 
-        room_id = self.helper.create_room_as(other_user, tok=other_user_token)
+        room_id = self.helper.create_room_as(self.other_user, tok=self.other_user_token)
 
         # Assert one user in room
         users_in_room = self.get_success(
             self.store.get_users_in_room(room_id),
         )
-        self.assertEqual([other_user], users_in_room)
+        self.assertEqual([self.other_user], users_in_room)
 
         # Enable require consent to send events
         self.event_creation_handler._block_events_without_consent_error = "Error"
-        self.event_creation_handler._consent_uri_builder = self._consent_uri_builder
 
         # Assert that the user is getting consent error
-        self.helper.send(room_id, body="foo", tok=other_user_token, expect_code=403)
-
-        # Mark the admin user as having consented
-        self.get_success(
-            self.store.user_set_consent_version(admin_user, "1"),
+        self.helper.send(
+            room_id,
+            body="foo", tok=self.other_user_token, expect_code=403,
         )
 
         # Test that the admin can still send shutdown
@@ -405,8 +410,8 @@ class ShutdownRoomTestCase(unittest.HomeserverTestCase):
         request, channel = self.make_request(
             "POST",
             url.encode('ascii'),
-            json.dumps({"new_room_user_id": admin_user}),
-            access_token=admin_user_tok,
+            json.dumps({"new_room_user_id": self.admin_user}),
+            access_token=self.admin_user_tok,
         )
         self.render(request)