From 8c2c7210efda3dd4d3480c25af1a23a63f152360 Mon Sep 17 00:00:00 2001
From: GB27247 <148772006+GB27247@users.noreply.github.com>
Date: Thu, 4 Jan 2024 09:24:40 +0000
Subject: [PATCH 1/2] Remove context manager and add testing

---
 lib/python-beta/src/bailo/core/client.py     |  7 ++----
 lib/python-beta/src/bailo/helper/model.py    |  2 +-
 lib/python-beta/src/bailo/helper/release.py  |  4 ++--
 lib/python-beta/tests/test_access_request.py |  3 +--
 lib/python-beta/tests/test_files.py          | 24 ++++++++++++++++++++
 5 files changed, 30 insertions(+), 10 deletions(-)
 create mode 100644 lib/python-beta/tests/test_files.py

diff --git a/lib/python-beta/src/bailo/core/client.py b/lib/python-beta/src/bailo/core/client.py
index 9c54446e9..b2954936c 100644
--- a/lib/python-beta/src/bailo/core/client.py
+++ b/lib/python-beta/src/bailo/core/client.py
@@ -311,11 +311,8 @@ def get_download_file(
         :param buffer: BytesIO object for bailo to write to
         :return: The unique file ID
         """
-        with self.agent.get(
-            f"{self.url}/v2/model/{model_id}/file/{file_id}/download", stream=True, timeout=10_000
-        ) as req:
-            with buffer as file:
-                shutil.copyfileobj(req.raw, file)
+        req = self.agent.get(f"{self.url}/v2/model/{model_id}/file/{file_id}/download", stream=True, timeout=10_000)
+        shutil.copyfileobj(req.raw, buffer)
         return file_id
 
     def simple_upload(self, model_id: str, name: str, buffer: BytesIO):
diff --git a/lib/python-beta/src/bailo/helper/model.py b/lib/python-beta/src/bailo/helper/model.py
index dd7fd2367..0e81e00d2 100644
--- a/lib/python-beta/src/bailo/helper/model.py
+++ b/lib/python-beta/src/bailo/helper/model.py
@@ -97,7 +97,7 @@ def update(self) -> None:
         )
         self.__unpack(res["model"])
 
-    def update_model_card(self, model_card: dict[str, Any] = None) -> None:
+    def update_model_card(self, model_card: dict[str, Any] | None = None) -> None:
         """Uploads and retrieves any changes to the model card on Bailo
 
         :param model_card: Model card dictionary, defaults to None
diff --git a/lib/python-beta/src/bailo/helper/release.py b/lib/python-beta/src/bailo/helper/release.py
index a9f7ff574..5fb09fac9 100644
--- a/lib/python-beta/src/bailo/helper/release.py
+++ b/lib/python-beta/src/bailo/helper/release.py
@@ -147,12 +147,12 @@ def upload(self, name: str, file: BytesIO) -> Any:
         :param name: The name of the file to upload to bailo
         :param f: A BytesIO object
 
-        :return: A JSON response object
+        :return: The unique file ID of the file uploaded
         """
         res = self.client.simple_upload(self.model_id, name, file).json()
         self.files.append(res["file"]["id"])
         self.update()
-        return res
+        return res["file"]["id"]
 
     def update(self) -> Any:
         """Updates the any changes to this release on Bailo.
diff --git a/lib/python-beta/tests/test_access_request.py b/lib/python-beta/tests/test_access_request.py
index b0f4e0631..baf17dafd 100644
--- a/lib/python-beta/tests/test_access_request.py
+++ b/lib/python-beta/tests/test_access_request.py
@@ -1,8 +1,7 @@
 from __future__ import annotations
 
 import pytest
-from bailo.core.client import Client
-from bailo.helper.access_request import AccessRequest
+from bailo import AccessRequest, Client
 
 
 def test_access_request():
diff --git a/lib/python-beta/tests/test_files.py b/lib/python-beta/tests/test_files.py
new file mode 100644
index 000000000..271b66343
--- /dev/null
+++ b/lib/python-beta/tests/test_files.py
@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from io import BytesIO
+
+import pytest
+from bailo import Model
+
+
+@pytest.mark.integration
+def test_file_upload(integration_client, example_model: Model):
+    byte_obj = b"Test Binary"
+    file = BytesIO(byte_obj)
+
+    example_release = example_model.create_release("0.1.0", "test")
+
+    with file as file:
+        file_id = example_release.upload("test", file)
+
+    download_file = BytesIO()
+    example_release.download(file_id, download_file)
+
+    # Check that file uploaded has the same contents as the one downloaded
+    download_file.seek(0)
+    assert download_file.read() == byte_obj

From d9905f68b5f0dcc1be783c88465ba275333796e6 Mon Sep 17 00:00:00 2001
From: GB27247 <148772006+GB27247@users.noreply.github.com>
Date: Thu, 4 Jan 2024 11:27:27 +0000
Subject: [PATCH 2/2] Bump version

---
 lib/python-beta/src/bailo/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/python-beta/src/bailo/__init__.py b/lib/python-beta/src/bailo/__init__.py
index 131a40c54..9d14f2254 100644
--- a/lib/python-beta/src/bailo/__init__.py
+++ b/lib/python-beta/src/bailo/__init__.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-__version__ = "0.2.2"
+__version__ = "0.2.3"
 
 from bailo.core.agent import Agent
 from bailo.core.client import Client