From 85a2e8662d94462c8f13f5c1a268b52e50ea49b6 Mon Sep 17 00:00:00 2001
From: Evan Blaudy <evan@cg-wire.com>
Date: Thu, 14 Mar 2024 15:20:54 +0100
Subject: [PATCH] [files] add functions to upload avatar for projects and
 persons

---
 gazu/files.py       | 33 +++++++++++++++++++++++++++++++--
 tests/test_files.py | 42 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 72 insertions(+), 3 deletions(-)

diff --git a/gazu/files.py b/gazu/files.py
index b36cfb48..edfbe9f6 100644
--- a/gazu/files.py
+++ b/gazu/files.py
@@ -1109,8 +1109,7 @@ def upload_working_file(working_file, file_path, client=default):
     """
     working_file = normalize_model_parameter(working_file)
     url_path = "/data/working-files/%s/file" % working_file["id"]
-    raw.upload(url_path, file_path, client=client)
-    return working_file
+    return raw.upload(url_path, file_path, client=client)
 
 
 def download_working_file(working_file, file_path=None, client=default):
@@ -1243,6 +1242,21 @@ def download_person_avatar(person, file_path, client=default):
     )
 
 
+def upload_person_avatar(person, file_path, client=default):
+    """
+    Upload given file as person avatar.
+
+    Args:
+        person (str / dict): The person dict or the person ID.
+        file_path (str): Path of the file to upload as avatar.
+    """
+    path = (
+        "/pictures/thumbnails/persons/%s"
+        % normalize_model_parameter(person)["id"]
+    )
+    return raw.upload(path, file_path, client=client)
+
+
 def download_project_avatar(project, file_path, client=default):
     """
     Download given project's avatar and save it at given location.
@@ -1258,6 +1272,21 @@ def download_project_avatar(project, file_path, client=default):
     )
 
 
+def upload_project_avatar(project, file_path, client=default):
+    """
+    Upload given file as project avatar.
+
+    Args:
+        project (str / dict): The project dict or the project ID.
+        file_path (str): Path of the file to upload as avatar.
+    """
+    path = (
+        "/pictures/thumbnails/projects/%s"
+        % normalize_model_parameter(project)["id"]
+    )
+    return raw.upload(path, file_path, client=client)
+
+
 def update_preview(preview_file, data, client=default):
     """
     Update the data of given preview file.
diff --git a/tests/test_files.py b/tests/test_files.py
index a8c585d7..6319db21 100644
--- a/tests/test_files.py
+++ b/tests/test_files.py
@@ -7,7 +7,7 @@
 import gazu.client
 import gazu.files
 
-from utils import fakeid, mock_route
+from utils import fakeid, mock_route, add_verify_file_callback
 
 
 class FilesTestCase(unittest.TestCase):
@@ -1170,3 +1170,43 @@ def test_get_output_file_by_path(self):
             self.assertEqual(
                 gazu.files.get_output_file_by_path("testpath"), text[0]
             )
+
+    def test_upload_person_avatar(self):
+        with open("./tests/fixtures/v1.png", "rb") as test_file:
+            with requests_mock.Mocker() as mock:
+                mock_route(
+                    mock,
+                    "POST",
+                    "pictures/thumbnails/persons/%s" % fakeid("person-1"),
+                    text={"id": fakeid("person-1")},
+                )
+
+                add_verify_file_callback(mock, {"file": test_file.read()})
+
+                self.assertEqual(
+                    gazu.files.upload_person_avatar(
+                        fakeid("person-1"),
+                        "./tests/fixtures/v1.png",
+                    ),
+                    {"id": fakeid("person-1")},
+                )
+
+    def test_upload_project_avatar(self):
+        with open("./tests/fixtures/v1.png", "rb") as test_file:
+            with requests_mock.Mocker() as mock:
+                mock_route(
+                    mock,
+                    "POST",
+                    "pictures/thumbnails/projects/%s" % fakeid("project-1"),
+                    text={"id": fakeid("project-1")},
+                )
+
+                add_verify_file_callback(mock, {"file": test_file.read()})
+
+                self.assertEqual(
+                    gazu.files.upload_project_avatar(
+                        fakeid("project-1"),
+                        "./tests/fixtures/v1.png",
+                    ),
+                    {"id": fakeid("project-1")},
+                )