Fix bug in upload file endpoint (#14924)

Co-authored-by: thomas chaton <[email protected]>
Co-authored-by: otaj <[email protected]>

Author: 3 people

src/lightning_app/CHANGELOG.md

@@ -27,6 +27,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - CLI usage without arguments errors ([#14877](https://github.com/Lightning-AI/lightning/pull/14877))
 
 
+- Fixed a bug where the upload files endpoint would raise an error when running locally ([#14924](https://github.com/Lightning-AI/lightning/pull/14924))
+
+
 ## [0.6.2] - 2022-09-21
 
 ### Changed

src/lightning_app/core/api.py

@@ -27,7 +27,8 @@
 from lightning_app.storage import Drive
 from lightning_app.utilities.app_helpers import InMemoryStateStore, Logger, StateStore
 from lightning_app.utilities.cloud import is_running_in_cloud
-from lightning_app.utilities.enum import OpenAPITags
+from lightning_app.utilities.component import _context
+from lightning_app.utilities.enum import ComponentContext, OpenAPITags
 from lightning_app.utilities.imports import _is_starsessions_available
 
 if _is_starsessions_available():
@@ -256,7 +257,8 @@ async def upload_file(filename: str, uploaded_file: UploadFile = File(...)):
                 f.write(content)
                 done = content == b""
 
-        drive.put(filename)
+        with _context(ComponentContext.WORK):
+            drive.put(filename)
     return f"Successfully uploaded '{filename}' to the Drive"
 
 

tests/tests_app/core/test_lightning_api.py

@@ -431,7 +431,7 @@ def target():
 
 
 def test_configure_api():
-
+    # Setup
     process = Process(target=target)
     process.start()
     time_left = 15
@@ -443,6 +443,13 @@ def test_configure_api():
             sleep(0.1)
             time_left -= 0.1
 
+    # Test Upload File
+    files = {"uploaded_file": open(__file__, "rb")}
+
+    response = requests.put(f"http://localhost:{APP_SERVER_PORT}/api/v1/upload_file/test", files=files)
+    assert response.json() == "Successfully uploaded 'test' to the Drive"
+
+    # Test Custom Request
     response = requests.post(
         f"http://localhost:{APP_SERVER_PORT}/api/v1/request", data=InputRequestModel(name="hello").json()
     )
@@ -451,6 +458,8 @@ def test_configure_api():
         f"http://localhost:{APP_SERVER_PORT}/api/v1/request", data=InputRequestModel(name="hello").json()
     )
     assert response.json() == {"name": "hello", "counter": 2}
+
+    # Teardown
     time_left = 15
     while time_left > 0:
         if process.exitcode == 0: