Add file upload, download, list, and delete operations

tests/conftest.py

@@ -12,6 +12,7 @@ from fastapi.testclient import TestClient
 from httpx import AsyncClient
 from langflow.graph.graph.base import Graph
 from langflow.services.auth.utils import get_password_hash
+from langflow.services.database.models.api_key.model import ApiKey
 from langflow.services.database.models.flow.model import Flow, FlowCreate
 from langflow.services.database.models.user.model import User, UserCreate
 from langflow.services.database.utils import session_getter
@@ -341,3 +342,22 @@ def added_vector_store(client, json_vector_store, logged_in_headers):
     assert response.json()["name"] == vector_store.name
     assert response.json()["data"] == vector_store.data
     return response.json()
+
+
+@pytest.fixture
+def created_api_key(active_user):
+    hashed = get_password_hash("random_key")
+    api_key = ApiKey(
+        name="test_api_key",
+        user_id=active_user.id,
+        api_key="random_key",
+        hashed_api_key=hashed,
+    )
+    db_manager = get_db_service()
+    with session_getter(db_manager) as session:
+        if existing_api_key := session.query(ApiKey).filter(ApiKey.api_key == api_key.api_key).first():
+            return existing_api_key
+        session.add(api_key)
+        session.commit()
+        session.refresh(api_key)
+    return api_key

tests/test_files.py

@@ -18,66 +18,73 @@ def mock_storage_service():
     return service
 
 
-def test_upload_file(client, mock_storage_service):
+def test_upload_file(client, mock_storage_service, created_api_key, flow):
+    headers = {"x-api-key": created_api_key.api_key}
     # Replace the actual storage service with the mock
     client.app.dependency_overrides[get_storage_service] = lambda: mock_storage_service
 
-    response = client.post("api/v1/files/upload/test_flow", files={"file": ("test.txt", b"test content")})
+    response = client.post(
+        f"api/v1/files/upload/{flow.id}", files={"file": ("test.txt", b"test content")}, headers=headers
+    )
     assert response.status_code == 201
-    assert response.json() == {"message": "File uploaded successfully", "file_path": "test_flow/test.txt"}
+    assert response.json() == {"flowId": str(flow.id), "file_path": f"{flow.id}/test.txt"}
 
 
-def test_download_file(client, mock_storage_service):
+def test_download_file(client, mock_storage_service, created_api_key, flow):
+    headers = {"x-api-key": created_api_key.api_key}
     client.app.dependency_overrides[get_storage_service] = lambda: mock_storage_service
 
-    response = client.get("api/v1/files/download/test_flow/test.txt")
+    response = client.get(f"api/v1/files/download/{flow.id}/test.txt", headers=headers)
     assert response.status_code == 200
     assert response.content == b"file content"
 
 
-def test_list_files(client, mock_storage_service):
+def test_list_files(client, mock_storage_service, created_api_key, flow):
+    headers = {"x-api-key": created_api_key.api_key}
     client.app.dependency_overrides[get_storage_service] = lambda: mock_storage_service
 
-    response = client.get("api/v1/files/list/test_flow")
+    response = client.get(f"api/v1/files/list/{flow.id}", headers=headers)
     assert response.status_code == 200
     assert response.json() == {"files": ["file1.txt", "file2.jpg"]}
 
 
-def test_delete_file(client, mock_storage_service):
+def test_delete_file(client, mock_storage_service, created_api_key, flow):
+    headers = {"x-api-key": created_api_key.api_key}
     client.app.dependency_overrides[get_storage_service] = lambda: mock_storage_service
 
-    response = client.delete("api/v1/files/delete/test_flow/test.txt")
+    response = client.delete(f"api/v1/files/delete/{flow.id}/test.txt", headers=headers)
     assert response.status_code == 200
     assert response.json() == {"message": "File test.txt deleted successfully"}
 
 
-def test_file_operations(client):
-    flow_id = "test_flow"
+def test_file_operations(client, created_api_key, flow):
+    headers = {"x-api-key": created_api_key.api_key}
+    flow_id = flow.id
     file_name = "test.txt"
     file_content = b"Hello, world!"
 
     # Step 1: Upload the file
-    response = client.post(f"api/v1/files/upload/{flow_id}", files={"file": (file_name, file_content)})
+    response = client.post(f"api/v1/files/upload/{flow_id}", files={"file": (file_name, file_content)}, headers=headers)
     assert response.status_code == 201
-    assert response.json() == {"message": "File uploaded successfully", "file_path": f"{flow_id}/{file_name}"}
+    assert response.json() == {"flowId": str(flow_id), "file_path": f"{flow_id}/{file_name}"}
 
     # Step 2: List files in the folder
-    response = client.get(f"api/v1/files/list/{flow_id}")
+    response = client.get(f"api/v1/files/list/{flow_id}", headers=headers)
     assert response.status_code == 200
     assert file_name in response.json()["files"]
 
     # Step 3: Download the file and verify its content
     mime_type = build_content_type_from_extension(file_name.split(".")[-1])
-    response = client.get(f"api/v1/files/download/{flow_id}/{file_name}")
+    response = client.get(f"api/v1/files/download/{flow_id}/{file_name}", headers=headers)
     assert response.status_code == 200
     assert response.content == file_content
     assert mime_type in response.headers["content-type"]
 
     # Step 4: Delete the file
-    response = client.delete(f"api/v1/files/delete/{flow_id}/{file_name}")
+    response = client.delete(f"api/v1/files/delete/{flow_id}/{file_name}", headers=headers)
     assert response.status_code == 200
     assert response.json() == {"message": f"File {file_name} deleted successfully"}
 
     # Verify that the file is indeed deleted
-    response = client.get(f"api/v1/files/list/{flow_id}")
+    response = client.get(f"api/v1/files/list/{flow_id}", headers=headers)
     assert file_name not in response.json()["files"]