Convert status codes to `HTTPStatus` in `tests.rest.admin` (#11455)

1 files changed, 33 insertions, 41 deletions

diff --git a/tests/rest/admin/test_admin.py b/tests/rest/admin/test_admin.py
index af849bd471..3adadcb46b 100644
--- a/tests/rest/admin/test_admin.py
+++ b/tests/rest/admin/test_admin.py
@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
 import os
 import urllib.parse
+from http import HTTPStatus
 from unittest.mock import Mock
 
 from twisted.internet.defer import Deferred
@@ -41,7 +41,7 @@ class VersionTestCase(unittest.HomeserverTestCase):
     def test_version_string(self):
         channel = self.make_request("GET", self.url, shorthand=False)
 
-        self.assertEqual(200, int(channel.result["code"]), msg=channel.result["body"])
+        self.assertEqual(HTTPStatus.OK, channel.code, msg=channel.json_body)
         self.assertEqual(
             {"server_version", "python_version"}, set(channel.json_body.keys())
         )
@@ -70,11 +70,11 @@ class DeleteGroupTestCase(unittest.HomeserverTestCase):
             content={"localpart": "test"},
         )
 
-        self.assertEqual(200, int(channel.result["code"]), msg=channel.result["body"])
+        self.assertEqual(HTTPStatus.OK, channel.code, msg=channel.json_body)
 
         group_id = channel.json_body["group_id"]
 
-        self._check_group(group_id, expect_code=200)
+        self._check_group(group_id, expect_code=HTTPStatus.OK)
 
         # Invite/join another user
 
@@ -82,13 +82,13 @@ class DeleteGroupTestCase(unittest.HomeserverTestCase):
         channel = self.make_request(
             "PUT", url.encode("ascii"), access_token=self.admin_user_tok, content={}
         )
-        self.assertEqual(200, int(channel.result["code"]), msg=channel.result["body"])
+        self.assertEqual(HTTPStatus.OK, channel.code, msg=channel.json_body)
 
         url = "/groups/%s/self/accept_invite" % (group_id,)
         channel = self.make_request(
             "PUT", url.encode("ascii"), access_token=self.other_user_token, content={}
         )
-        self.assertEqual(200, int(channel.result["code"]), msg=channel.result["body"])
+        self.assertEqual(HTTPStatus.OK, channel.code, msg=channel.json_body)
 
         # Check other user knows they're in the group
         self.assertIn(group_id, self._get_groups_user_is_in(self.admin_user_tok))
@@ -103,10 +103,10 @@ class DeleteGroupTestCase(unittest.HomeserverTestCase):
             content={"localpart": "test"},
         )
 
-        self.assertEqual(200, int(channel.result["code"]), msg=channel.result["body"])
+        self.assertEqual(HTTPStatus.OK, channel.code, msg=channel.json_body)
 
-        # Check group returns 404
-        self._check_group(group_id, expect_code=404)
+        # Check group returns HTTPStatus.NOT_FOUND
+        self._check_group(group_id, expect_code=HTTPStatus.NOT_FOUND)
 
         # Check users don't think they're in the group
         self.assertNotIn(group_id, self._get_groups_user_is_in(self.admin_user_tok))
@@ -122,15 +122,13 @@ class DeleteGroupTestCase(unittest.HomeserverTestCase):
             "GET", url.encode("ascii"), access_token=self.admin_user_tok
         )
 
-        self.assertEqual(
-            expect_code, int(channel.result["code"]), msg=channel.result["body"]
-        )
+        self.assertEqual(expect_code, channel.code, msg=channel.json_body)
 
     def _get_groups_user_is_in(self, access_token):
         """Returns the list of groups the user is in (given their access token)"""
         channel = self.make_request("GET", b"/joined_groups", access_token=access_token)
 
-        self.assertEqual(200, int(channel.result["code"]), msg=channel.result["body"])
+        self.assertEqual(HTTPStatus.OK, channel.code, msg=channel.json_body)
 
         return channel.json_body["groups"]
 
@@ -210,10 +208,10 @@ class QuarantineMediaTestCase(unittest.HomeserverTestCase):
 
         # Should be quarantined
         self.assertEqual(
-            404,
-            int(channel.code),
+            HTTPStatus.NOT_FOUND,
+            channel.code,
             msg=(
-                "Expected to receive a 404 on accessing quarantined media: %s"
+                "Expected to receive a HTTPStatus.NOT_FOUND on accessing quarantined media: %s"
                 % server_and_media_id
             ),
         )
@@ -232,8 +230,8 @@ class QuarantineMediaTestCase(unittest.HomeserverTestCase):
 
         # Expect a forbidden error
         self.assertEqual(
-            403,
-            int(channel.result["code"]),
+            HTTPStatus.FORBIDDEN,
+            channel.code,
             msg="Expected forbidden on quarantining media as a non-admin",
         )
 
@@ -247,8 +245,8 @@ class QuarantineMediaTestCase(unittest.HomeserverTestCase):
 
         # Expect a forbidden error
         self.assertEqual(
-            403,
-            int(channel.result["code"]),
+            HTTPStatus.FORBIDDEN,
+            channel.code,
             msg="Expected forbidden on quarantining media as a non-admin",
         )
 
@@ -279,7 +277,7 @@ class QuarantineMediaTestCase(unittest.HomeserverTestCase):
         )
 
         # Should be successful
-        self.assertEqual(200, int(channel.code), msg=channel.result["body"])
+        self.assertEqual(HTTPStatus.OK, channel.code)
 
         # Quarantine the media
         url = "/_synapse/admin/v1/media/quarantine/%s/%s" % (
@@ -292,7 +290,7 @@ class QuarantineMediaTestCase(unittest.HomeserverTestCase):
             access_token=admin_user_tok,
         )
         self.pump(1.0)
-        self.assertEqual(200, int(channel.code), msg=channel.result["body"])
+        self.assertEqual(HTTPStatus.OK, channel.code, msg=channel.json_body)
 
         # Attempt to access the media
         self._ensure_quarantined(admin_user_tok, server_name_and_media_id)
@@ -348,11 +346,9 @@ class QuarantineMediaTestCase(unittest.HomeserverTestCase):
             access_token=admin_user_tok,
         )
         self.pump(1.0)
-        self.assertEqual(200, int(channel.code), msg=channel.result["body"])
+        self.assertEqual(HTTPStatus.OK, channel.code, msg=channel.json_body)
         self.assertEqual(
-            json.loads(channel.result["body"].decode("utf-8")),
-            {"num_quarantined": 2},
-            "Expected 2 quarantined items",
+            channel.json_body, {"num_quarantined": 2}, "Expected 2 quarantined items"
         )
 
         # Convert mxc URLs to server/media_id strings
@@ -396,11 +392,9 @@ class QuarantineMediaTestCase(unittest.HomeserverTestCase):
             access_token=admin_user_tok,
         )
         self.pump(1.0)
-        self.assertEqual(200, int(channel.result["code"]), msg=channel.result["body"])
+        self.assertEqual(HTTPStatus.OK, channel.code, msg=channel.json_body)
         self.assertEqual(
-            json.loads(channel.result["body"].decode("utf-8")),
-            {"num_quarantined": 2},
-            "Expected 2 quarantined items",
+            channel.json_body, {"num_quarantined": 2}, "Expected 2 quarantined items"
         )
 
         # Attempt to access each piece of media
@@ -432,7 +426,7 @@ class QuarantineMediaTestCase(unittest.HomeserverTestCase):
         url = "/_synapse/admin/v1/media/protect/%s" % (urllib.parse.quote(media_id_2),)
         channel = self.make_request("POST", url, access_token=admin_user_tok)
         self.pump(1.0)
-        self.assertEqual(200, int(channel.code), msg=channel.result["body"])
+        self.assertEqual(HTTPStatus.OK, channel.code, msg=channel.json_body)
 
         # Quarantine all media by this user
         url = "/_synapse/admin/v1/user/%s/media/quarantine" % urllib.parse.quote(
@@ -444,11 +438,9 @@ class QuarantineMediaTestCase(unittest.HomeserverTestCase):
             access_token=admin_user_tok,
         )
         self.pump(1.0)
-        self.assertEqual(200, int(channel.result["code"]), msg=channel.result["body"])
+        self.assertEqual(HTTPStatus.OK, channel.code, msg=channel.json_body)
         self.assertEqual(
-            json.loads(channel.result["body"].decode("utf-8")),
-            {"num_quarantined": 1},
-            "Expected 1 quarantined item",
+            channel.json_body, {"num_quarantined": 1}, "Expected 1 quarantined item"
         )
 
         # Attempt to access each piece of media, the first should fail, the
@@ -467,10 +459,10 @@ class QuarantineMediaTestCase(unittest.HomeserverTestCase):
 
         # Shouldn't be quarantined
         self.assertEqual(
-            200,
-            int(channel.code),
+            HTTPStatus.OK,
+            channel.code,
             msg=(
-                "Expected to receive a 200 on accessing not-quarantined media: %s"
+                "Expected to receive a HTTPStatus.OK on accessing not-quarantined media: %s"
                 % server_and_media_id_2
             ),
         )
@@ -499,7 +491,7 @@ class PurgeHistoryTestCase(unittest.HomeserverTestCase):
     def test_purge_history(self):
         """
         Simple test of purge history API.
-        Test only that is is possible to call, get status 200 and purge_id.
+        Test only that is is possible to call, get status HTTPStatus.OK and purge_id.
         """
 
         channel = self.make_request(
@@ -509,7 +501,7 @@ class PurgeHistoryTestCase(unittest.HomeserverTestCase):
             access_token=self.admin_user_tok,
         )
 
-        self.assertEqual(200, channel.code, msg=channel.json_body)
+        self.assertEqual(HTTPStatus.OK, channel.code, msg=channel.json_body)
         self.assertIn("purge_id", channel.json_body)
         purge_id = channel.json_body["purge_id"]
 
@@ -520,5 +512,5 @@ class PurgeHistoryTestCase(unittest.HomeserverTestCase):
             access_token=self.admin_user_tok,
         )
 
-        self.assertEqual(200, channel.code, msg=channel.json_body)
+        self.assertEqual(HTTPStatus.OK, channel.code, msg=channel.json_body)
         self.assertEqual("complete", channel.json_body["status"])