test_servlet.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138
  1. # Copyright 2020 The Matrix.org Foundation C.I.C.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import json
  15. from http import HTTPStatus
  16. from io import BytesIO
  17. from typing import Tuple
  18. from unittest.mock import Mock
  19. from synapse.api.errors import Codes, SynapseError
  20. from synapse.http.servlet import (
  21. RestServlet,
  22. parse_json_object_from_request,
  23. parse_json_value_from_request,
  24. )
  25. from synapse.http.site import SynapseRequest
  26. from synapse.rest.client._base import client_patterns
  27. from synapse.server import HomeServer
  28. from synapse.types import JsonDict
  29. from synapse.util.cancellation import cancellable
  30. from tests import unittest
  31. from tests.http.server._base import test_disconnect
  32. def make_request(content):
  33. """Make an object that acts enough like a request."""
  34. request = Mock(spec=["method", "uri", "content"])
  35. if isinstance(content, dict):
  36. content = json.dumps(content).encode("utf8")
  37. request.method = bytes("STUB_METHOD", "ascii")
  38. request.uri = bytes("/test_stub_uri", "ascii")
  39. request.content = BytesIO(content)
  40. return request
  41. class TestServletUtils(unittest.TestCase):
  42. def test_parse_json_value(self):
  43. """Basic tests for parse_json_value_from_request."""
  44. # Test round-tripping.
  45. obj = {"foo": 1}
  46. result1 = parse_json_value_from_request(make_request(obj))
  47. self.assertEqual(result1, obj)
  48. # Results don't have to be objects.
  49. result2 = parse_json_value_from_request(make_request(b'["foo"]'))
  50. self.assertEqual(result2, ["foo"])
  51. # Test empty.
  52. with self.assertRaises(SynapseError):
  53. parse_json_value_from_request(make_request(b""))
  54. result3 = parse_json_value_from_request(
  55. make_request(b""), allow_empty_body=True
  56. )
  57. self.assertIsNone(result3)
  58. # Invalid UTF-8.
  59. with self.assertRaises(SynapseError):
  60. parse_json_value_from_request(make_request(b"\xFF\x00"))
  61. # Invalid JSON.
  62. with self.assertRaises(SynapseError):
  63. parse_json_value_from_request(make_request(b"foo"))
  64. with self.assertRaises(SynapseError):
  65. parse_json_value_from_request(make_request(b'{"foo": Infinity}'))
  66. def test_parse_json_object(self):
  67. """Basic tests for parse_json_object_from_request."""
  68. # Test empty.
  69. result = parse_json_object_from_request(
  70. make_request(b""), allow_empty_body=True
  71. )
  72. self.assertEqual(result, {})
  73. # Test not an object
  74. with self.assertRaises(SynapseError):
  75. parse_json_object_from_request(make_request(b'["foo"]'))
  76. class CancellableRestServlet(RestServlet):
  77. """A `RestServlet` with a mix of cancellable and uncancellable handlers."""
  78. PATTERNS = client_patterns("/sleep$")
  79. def __init__(self, hs: HomeServer):
  80. super().__init__()
  81. self.clock = hs.get_clock()
  82. @cancellable
  83. async def on_GET(self, request: SynapseRequest) -> Tuple[int, JsonDict]:
  84. await self.clock.sleep(1.0)
  85. return HTTPStatus.OK, {"result": True}
  86. async def on_POST(self, request: SynapseRequest) -> Tuple[int, JsonDict]:
  87. await self.clock.sleep(1.0)
  88. return HTTPStatus.OK, {"result": True}
  89. class TestRestServletCancellation(unittest.HomeserverTestCase):
  90. """Tests for `RestServlet` cancellation."""
  91. servlets = [
  92. lambda hs, http_server: CancellableRestServlet(hs).register(http_server)
  93. ]
  94. def test_cancellable_disconnect(self) -> None:
  95. """Test that handlers with the `@cancellable` flag can be cancelled."""
  96. channel = self.make_request("GET", "/sleep", await_result=False)
  97. test_disconnect(
  98. self.reactor,
  99. channel,
  100. expect_cancellation=True,
  101. expected_body={"error": "Request cancelled", "errcode": Codes.UNKNOWN},
  102. )
  103. def test_uncancellable_disconnect(self) -> None:
  104. """Test that handlers without the `@cancellable` flag cannot be cancelled."""
  105. channel = self.make_request("POST", "/sleep", await_result=False)
  106. test_disconnect(
  107. self.reactor,
  108. channel,
  109. expect_cancellation=False,
  110. expected_body={"result": True},
  111. )