From 9968a6b1345f6ea2ea42e1aa4c5e160565afe4aa Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Sun, 28 Jan 2024 12:01:08 -0300 Subject: [PATCH] Add file upload, download, list, and delete operations --- tests/conftest.py | 20 ++++++++++++++++++++ tests/test_files.py | 41 ++++++++++++++++++++++++----------------- 2 files changed, 44 insertions(+), 17 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 6538ddbe1..309185a74 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/test_files.py b/tests/test_files.py index 45d9c1ebb..ad66f1497 100644 --- a/tests/test_files.py +++ b/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"]