From d2b860c0c8b06c7a8ef35fa57b69e5a702288481 Mon Sep 17 00:00:00 2001
From: Yusuf Olokoba <olokobayusuf@gmail.com>
Date: Tue, 20 Feb 2024 13:52:58 -0500
Subject: [PATCH] Simplify URLs in `fxn.storage.upload`

---
 fxn/services/storage.py | 13 +++++++++++--
 test/storage_test.py    |  6 +++---
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/fxn/services/storage.py b/fxn/services/storage.py
index 8f55881..5a59ba6 100644
--- a/fxn/services/storage.py
+++ b/fxn/services/storage.py
@@ -10,6 +10,7 @@
 from requests import put
 from rich.progress import open as open_progress, wrap_file
 from typing import Union
+from urllib.parse import urlparse, urlunparse
 
 from ..graph import GraphClient
 from ..types import UploadType
@@ -97,7 +98,7 @@ def __upload_file (
         with open_progress(file, mode="rb", description=name, disable=not verbose) as f:
             put(url, data=f, headers={ "Content-Type": mime }).raise_for_status()
         # Return
-        return url
+        return self.__simplify_url(url)
     
     def __upload_buffer (
         self,
@@ -122,9 +123,17 @@ def __upload_buffer (
         with wrap_file(file, total=size, description=name, disable=not verbose) as f:
             put(url, data=f, headers={ "Content-Type": mime }).raise_for_status()
         # Return
-        return url
+        return self.__simplify_url(url)
     
     def __create_data_url (self, file: BytesIO, *, mime: str) -> str:
         encoded_data = b64encode(file.getvalue()).decode("ascii")
         url = f"data:{mime};base64,{encoded_data}"
+        return url
+    
+    def __simplify_url (self, url: str) -> str:
+        if url.startswith("data:"):
+            return url
+        parsed_url = urlparse(url)
+        parsed_url = parsed_url._replace(netloc="cdn.fxn.ai", query="")
+        url = urlunparse(parsed_url)
         return url
\ No newline at end of file
diff --git a/test/storage_test.py b/test/storage_test.py
index 24c74df..b79c7ff 100644
--- a/test/storage_test.py
+++ b/test/storage_test.py
@@ -9,8 +9,8 @@
 
 def test_file_upload ():
     fxn = Function()
-    url = fxn.storage.upload("test/media/cat.jpg", UploadType.Value, verbose=True)
-    assert url.startswith("https://")
+    url = fxn.storage.upload("test/media/cat.jpg", type=UploadType.Value, verbose=True)
+    assert url.startswith("https://cdn.fxn.ai")
 
 def test_file_upload_data_url ():
     fxn = Function()
@@ -25,7 +25,7 @@ def test_buffer_upload ():
     with open(path, mode="rb") as f:
         buffer = BytesIO(f.read())
     url = fxn.storage.upload(buffer, type=UploadType.Value, name=path.name, verbose=True)
-    assert url.startswith("https://")
+    assert url.startswith("https://cdn.fxn.ai")
 
 def test_buffer_upload_data_url ():
     fxn = Function()