From 3d10fa472baafc4995b5f8d75de95ce3794e400a Mon Sep 17 00:00:00 2001
From: Evgeniy Alekseev <i@arcanis.me>
Date: Wed, 18 Aug 2021 05:04:26 +0300
Subject: [PATCH] guess mime type for local files

---
 src/ahriman/core/upload/s3.py        |  8 +++++++-
 tests/ahriman/core/upload/test_s3.py | 24 +++++++++++++++++-------
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/src/ahriman/core/upload/s3.py b/src/ahriman/core/upload/s3.py
index 008d6831..b5182da3 100644
--- a/src/ahriman/core/upload/s3.py
+++ b/src/ahriman/core/upload/s3.py
@@ -19,6 +19,7 @@
 #
 import boto3  # type: ignore
 import hashlib
+import mimetypes
 
 from pathlib import Path
 from typing import Any, Dict, Generator, Iterable
@@ -122,8 +123,13 @@ class S3(Upload):
             remote_checksum = remote_object.e_tag[1:-1] if remote_object is not None else None
             if remote_checksum == checksum:
                 continue
+
+            local_path = path / local_file
             remote_path = Path(self.architecture) / local_file
-            self.bucket.upload_file(str(path / local_file), str(remote_path))
+            (mime, _) = mimetypes.guess_type(local_path)
+            extra_args = {"Content-Type": mime} if mime is not None else None
+
+            self.bucket.upload_file(Filename=str(local_path), Key=str(remote_path), ExtraArgs=extra_args)
 
         # remove files which were removed locally
         for local_file, remote_object in remote_objects.items():
diff --git a/tests/ahriman/core/upload/test_s3.py b/tests/ahriman/core/upload/test_s3.py
index 04e11631..288dffbd 100644
--- a/tests/ahriman/core/upload/test_s3.py
+++ b/tests/ahriman/core/upload/test_s3.py
@@ -1,6 +1,6 @@
 from pathlib import Path
 from pytest_mock import MockerFixture
-from typing import Any, List
+from typing import Any, List, Optional, Tuple
 from unittest import mock
 from unittest.mock import MagicMock
 
@@ -74,15 +74,17 @@ def test_sync(s3: S3, s3_remote_objects: List[Any], mocker: MockerFixture) -> No
     """
     must run sync command
     """
+    def mimetype(path: Path) -> Tuple[Optional[str], None]:
+        return ("text/html", None) if path.name == "b" else (None, None)
+
     root = Path("path")
     local_files = {
         Path(item.key.replace("a", "d")): item.e_tag.replace("b", "d").replace("\"", "")
         for item in s3_remote_objects
     }
     remote_objects = {Path(item.key): item for item in s3_remote_objects}
-    print(local_files)
-    print(remote_objects)
 
+    mocker.patch("mimetypes.guess_type", side_effect=mimetype)
     local_files_mock = mocker.patch("ahriman.core.upload.s3.S3.get_local_files", return_value=local_files)
     remote_objects_mock = mocker.patch("ahriman.core.upload.s3.S3.get_remote_objects", return_value=remote_objects)
     upload_mock = s3.bucket = MagicMock()
@@ -91,8 +93,16 @@ def test_sync(s3: S3, s3_remote_objects: List[Any], mocker: MockerFixture) -> No
 
     local_files_mock.assert_called_once()
     remote_objects_mock.assert_called_once()
-    upload_mock.upload_file.assert_has_calls([
-        mock.call(str(root / s3.architecture / "b"), f"{s3.architecture}/{s3.architecture}/b"),
-        mock.call(str(root / s3.architecture / "d"), f"{s3.architecture}/{s3.architecture}/d"),
-    ], any_order=True)
+    upload_mock.upload_file.assert_has_calls(
+        [
+            mock.call(
+                Filename=str(root / s3.architecture / "b"),
+                Key=f"{s3.architecture}/{s3.architecture}/b",
+                ExtraArgs={"Content-Type": "text/html"}),
+            mock.call(
+                Filename=str(root / s3.architecture / "d"),
+                Key=f"{s3.architecture}/{s3.architecture}/d",
+                ExtraArgs=None),
+        ],
+        any_order=True)
     remote_objects[Path("x86_64/a")].delete.assert_called_once()