2 files changed, 65 insertions, 4 deletions
diff --git a/tests/media/test_base.py b/tests/media/test_base.py
index 4728c80969..119d7ba66f 100644
--- a/tests/media/test_base.py
+++ b/tests/media/test_base.py
@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from synapse.media._base import get_filename_from_headers
+from unittest.mock import Mock
+
+from synapse.media._base import add_file_headers, get_filename_from_headers
 
 from tests import unittest
 
@@ -36,3 +38,28 @@ class GetFileNameFromHeadersTests(unittest.TestCase):
                 expected,
                 f"expected output for {hdr!r} to be {expected} but was {res}",
             )
+
+
+class AddFileHeadersTests(unittest.TestCase):
+    TEST_CASES = {
+        "text/plain": b"inline; filename=file.name",
+        "text/csv": b"inline; filename=file.name",
+        "image/png": b"inline; filename=file.name",
+        "text/html": b"attachment; filename=file.name",
+        "any/thing": b"attachment; filename=file.name",
+    }
+
+    def test_content_disposition(self) -> None:
+        for media_type, expected in self.TEST_CASES.items():
+            request = Mock()
+            add_file_headers(request, media_type, 0, "file.name")
+            request.setHeader.assert_any_call(b"Content-Disposition", expected)
+
+    def test_no_filename(self) -> None:
+        request = Mock()
+        add_file_headers(request, "text/plain", 0, None)
+        request.setHeader.assert_any_call(b"Content-Disposition", b"inline")
+
+        request.reset_mock()
+        add_file_headers(request, "text/html", 0, None)
+        request.setHeader.assert_any_call(b"Content-Disposition", b"attachment")
diff --git a/tests/media/test_media_storage.py b/tests/media/test_media_storage.py
index ea0051dde4..04fc7bdcef 100644
--- a/tests/media/test_media_storage.py
+++ b/tests/media/test_media_storage.py
@@ -129,6 +129,8 @@ class _TestImage:
             a 404/400 is expected.
         unable_to_thumbnail: True if we expect the thumbnailing to fail (400), or
             False if the thumbnailing should succeed or a normal 404 is expected.
+        is_inline: True if we expect the file to be served using an inline
+            Content-Disposition or False if we expect an attachment.
     """
 
     data: bytes
@@ -138,6 +140,7 @@ class _TestImage:
     expected_scaled: Optional[bytes] = None
     expected_found: bool = True
     unable_to_thumbnail: bool = False
+    is_inline: bool = True
 
 
 @parameterized_class(
@@ -198,6 +201,25 @@ class _TestImage:
                 unable_to_thumbnail=True,
             ),
         ),
+        # An SVG.
+        (
+            _TestImage(
+                b"""<?xml version="1.0"?>
+<!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN"
+  "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">
+
+<svg xmlns="http://www.w3.org/2000/svg"
+     width="400" height="400">
+  <circle cx="100" cy="100" r="50" stroke="black"
+    stroke-width="5" fill="red" />
+</svg>""",
+                b"image/svg",
+                b".svg",
+                expected_found=False,
+                unable_to_thumbnail=True,
+                is_inline=False,
+            ),
+        ),
     ],
 )
 class MediaRepoTests(unittest.HomeserverTestCase):
@@ -339,7 +361,11 @@ class MediaRepoTests(unittest.HomeserverTestCase):
         )
         self.assertEqual(
             headers.getRawHeaders(b"Content-Disposition"),
-            [b"attachment; filename=out" + self.test_image.extension],
+            [
+                (b"inline" if self.test_image.is_inline else b"attachment")
+                + b"; filename=out"
+                + self.test_image.extension
+            ],
         )
 
     def test_disposition_filenamestar_utf8escaped(self) -> None:
@@ -359,7 +385,12 @@ class MediaRepoTests(unittest.HomeserverTestCase):
         )
         self.assertEqual(
             headers.getRawHeaders(b"Content-Disposition"),
-            [b"attachment; filename*=utf-8''" + filename + self.test_image.extension],
+            [
+                (b"inline" if self.test_image.is_inline else b"attachment")
+                + b"; filename*=utf-8''"
+                + filename
+                + self.test_image.extension
+            ],
         )
 
     def test_disposition_none(self) -> None:
@@ -373,7 +404,10 @@ class MediaRepoTests(unittest.HomeserverTestCase):
         self.assertEqual(
             headers.getRawHeaders(b"Content-Type"), [self.test_image.content_type]
         )
-        self.assertEqual(headers.getRawHeaders(b"Content-Disposition"), [b"attachment"])
+        self.assertEqual(
+            headers.getRawHeaders(b"Content-Disposition"),
+            [b"inline" if self.test_image.is_inline else b"attachment"],
+        )
 
     def test_thumbnail_crop(self) -> None:
         """Test that a cropped remote thumbnail is available."""