Fix bug in upload file endpoint

src/lightning_app/core/api.py

@@ -26,7 +26,8 @@
 from lightning_app.core.queues import RedisQueue
 from lightning_app.storage import Drive
 from lightning_app.utilities.app_helpers import InMemoryStateStore, Logger, StateStore
-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_redis_available, _is_starsessions_available
 
 if _is_starsessions_available():
@@ -255,7 +256,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

@@ -430,7 +430,7 @@ def target():
 
 
 def test_configure_api():
-
+    # Setup
     process = Process(target=target)
     process.start()
     time_left = 15
@@ -442,6 +442,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()
     )
@@ -450,6 +457,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: