|
| 1 | +# |
| 2 | +# This file is licensed under the Affero General Public License (AGPL) version 3. |
| 3 | +# |
| 4 | +# Copyright (C) 2026 Element Creations Ltd |
| 5 | +# |
| 6 | +# This program is free software: you can redistribute it and/or modify |
| 7 | +# it under the terms of the GNU Affero General Public License as |
| 8 | +# published by the Free Software Foundation, either version 3 of the |
| 9 | +# License, or (at your option) any later version. |
| 10 | +# |
| 11 | +# See the GNU Affero General Public License for more details: |
| 12 | +# <https://www.gnu.org/licenses/agpl-3.0.html>. |
| 13 | +# |
| 14 | + |
| 15 | +from synapse.replication.tcp.streams import QuarantinedMediaStream |
| 16 | +from synapse.types import UserID |
| 17 | + |
| 18 | +from tests.replication._base import BaseMultiWorkerStreamTestCase |
| 19 | + |
| 20 | + |
| 21 | +class QuarantinedMediaWorkerWriterTestCase(BaseMultiWorkerStreamTestCase): |
| 22 | + """Checks that the quarantined_media stream is replicated when the |
| 23 | + configured stream writer is a worker rather than the main process. |
| 24 | + """ |
| 25 | + |
| 26 | + def default_config(self) -> dict: |
| 27 | + conf = super().default_config() |
| 28 | + conf["stream_writers"] = {"quarantined_media_changes": ["worker1"]} |
| 29 | + conf["instance_map"] = { |
| 30 | + "main": {"host": "testserv", "port": 8765}, |
| 31 | + "worker1": {"host": "testserv", "port": 1001}, |
| 32 | + } |
| 33 | + return conf |
| 34 | + |
| 35 | + def test_quarantine_on_worker_writer_replicates_to_main(self) -> None: |
| 36 | + main_store = self.hs.get_datastores().main |
| 37 | + |
| 38 | + worker_hs = self.make_worker_hs( |
| 39 | + "synapse.app.generic_worker", {"worker_name": "worker1"} |
| 40 | + ) |
| 41 | + worker_store = worker_hs.get_datastores().main |
| 42 | + |
| 43 | + # The worker must consider itself a source of the stream... |
| 44 | + self.assertIn( |
| 45 | + QuarantinedMediaStream.NAME, |
| 46 | + { |
| 47 | + stream.NAME |
| 48 | + for stream in worker_hs.get_replication_command_handler().get_streams_to_replicate() |
| 49 | + }, |
| 50 | + ) |
| 51 | + # ... and the main process must not, as it isn't a writer. |
| 52 | + self.assertNotIn( |
| 53 | + QuarantinedMediaStream.NAME, |
| 54 | + { |
| 55 | + stream.NAME |
| 56 | + for stream in self.hs.get_replication_command_handler().get_streams_to_replicate() |
| 57 | + }, |
| 58 | + ) |
| 59 | + |
| 60 | + # Quarantining only records a change for media that exists. |
| 61 | + self.get_success( |
| 62 | + main_store.store_local_media( |
| 63 | + media_id="media_id1", |
| 64 | + media_type="text/plain", |
| 65 | + time_now_ms=self.clock.time_msec(), |
| 66 | + upload_name=None, |
| 67 | + media_length=100, |
| 68 | + user_id=UserID.from_string("@user:test"), |
| 69 | + ) |
| 70 | + ) |
| 71 | + |
| 72 | + initial_token = main_store.get_current_quarantined_media_stream_id() |
| 73 | + |
| 74 | + # Quarantine the media on the worker, i.e. the configured writer. |
| 75 | + self.get_success( |
| 76 | + worker_store.quarantine_media_by_id("test", "media_id1", "@admin:test") |
| 77 | + ) |
| 78 | + |
| 79 | + self.replicate() |
| 80 | + |
| 81 | + # The main process only learns of the new stream ID over replication, |
| 82 | + # even though the two instances share a database. |
| 83 | + self.assertEqual( |
| 84 | + main_store.get_current_quarantined_media_stream_id(), |
| 85 | + initial_token + 1, |
| 86 | + ) |
0 commit comments