Test cases

Author: hughns

tests/media/test_media_storage.py

@@ -1362,3 +1362,42 @@ async def _send_request(*args: Any, **kwargs: Any) -> IResponse:
             store_media.sha256,
             SMALL_PNG_SHA256,
         )
+
+
+class MediaRepoSizeModuleCallbackTestCase(unittest.HomeserverTestCase):
+    servlets = [
+        login.register_servlets,
+        admin.register_servlets,
+    ]
+
+    def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
+        self.user = self.register_user("user", "pass")
+        self.tok = self.login("user", "pass")
+        self.mock_result = True  # Allow all uploads by default
+
+        hs.get_module_api().register_media_repository_callbacks(
+            is_user_allowed_to_upload_media_of_size=self.is_user_allowed_to_upload_media_of_size,
+        )
+
+    def create_resource_dict(self) -> Dict[str, Resource]:
+        resources = super().create_resource_dict()
+        resources["/_matrix/media"] = self.hs.get_media_repository_resource()
+        return resources
+
+    async def is_user_allowed_to_upload_media_of_size(
+        self, user_id: str, size: int
+    ) -> bool:
+        self.last_user_id = user_id
+        self.last_size = size
+        return self.mock_result
+
+    def test_upload_allowed(self) -> None:
+        self.helper.upload_media(SMALL_PNG, tok=self.tok, expect_code=200)
+        assert self.last_user_id == self.user
+        assert self.last_size == len(SMALL_PNG)
+
+    def test_upload_not_allowed(self) -> None:
+        self.mock_result = False
+        self.helper.upload_media(SMALL_PNG, tok=self.tok, expect_code=413)
+        assert self.last_user_id == self.user
+        assert self.last_size == len(SMALL_PNG)

tests/rest/client/test_media.py

@@ -1618,6 +1618,63 @@ def test_media_config(self) -> None:
         )
 
 
+class MediaConfigModuleCallbackTestCase(unittest.HomeserverTestCase):
+    servlets = [
+        media.register_servlets,
+        admin.register_servlets,
+        login.register_servlets,
+    ]
+
+    def make_homeserver(
+        self, reactor: ThreadedMemoryReactorClock, clock: Clock
+    ) -> HomeServer:
+        config = self.default_config()
+
+        self.storage_path = self.mktemp()
+        self.media_store_path = self.mktemp()
+        os.mkdir(self.storage_path)
+        os.mkdir(self.media_store_path)
+        config["media_store_path"] = self.media_store_path
+
+        provider_config = {
+            "module": "synapse.media.storage_provider.FileStorageProviderBackend",
+            "store_local": True,
+            "store_synchronous": False,
+            "store_remote": True,
+            "config": {"directory": self.storage_path},
+        }
+
+        config["media_storage_providers"] = [provider_config]
+
+        return self.setup_test_homeserver(config=config)
+
+    def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
+        self.user = self.register_user("user", "password")
+        self.tok = self.login("user", "password")
+
+        hs.get_module_api().register_media_repository_callbacks(
+            get_media_config_for_user=self.get_media_config_for_user,
+        )
+
+    async def get_media_config_for_user(
+        self,
+        user_id: str,
+    ) -> Optional[JsonDict]:
+        # We echo back the user_id and set a custom upload size.
+        return {"m.upload.size": 1024, "user_id": user_id}
+
+    def test_media_config(self) -> None:
+        channel = self.make_request(
+            "GET",
+            "/_matrix/client/v1/media/config",
+            shorthand=False,
+            access_token=self.tok,
+        )
+        self.assertEqual(channel.code, 200)
+        self.assertEqual(channel.json_body["m.upload.size"], 1024)
+        self.assertEqual(channel.json_body["user_id"], self.user)
+
+
 class RemoteDownloadLimiterTestCase(unittest.HomeserverTestCase):
     servlets = [
         media.register_servlets,