adds vector store tests

This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-09-17 19:34:46 -03:00
commit 6871cc5ab6
3 changed files with 1357 additions and 1 deletions

View file

@ -33,7 +33,9 @@ def pytest_configure():
pytest.BASIC_CHAT_WITH_PROMPT_AND_HISTORY = (
Path(__file__).parent.absolute() / "data" / "BasicChatwithPromptandHistory.json"
)
pytest.VECTOR_STORE_PATH = (
Path(__file__).parent.absolute() / "data" / "Vector_store.json"
)
pytest.CODE_WITH_SYNTAX_ERROR = """
def get_text():
retun "Hello World"
@ -111,6 +113,12 @@ def json_flow_with_prompt_and_history():
return f.read()
@pytest.fixture
def json_vector_store():
with open(pytest.VECTOR_STORE_PATH, "r") as f:
return f.read()
@pytest.fixture(name="session")
def session_fixture():
engine = create_engine(
@ -230,3 +238,17 @@ def added_flow(client, json_flow_with_prompt_and_history, logged_in_headers):
assert response.json()["name"] == flow.name
assert response.json()["data"] == flow.data
return response.json()
@pytest.fixture
def added_vector_store(client, json_vector_store, logged_in_headers):
vector_store = orjson.loads(json_vector_store)
data = vector_store["data"]
vector_store = FlowCreate(name="Vector Store", description="description", data=data)
response = client.post(
"api/v1/flows/", json=vector_store.dict(), headers=logged_in_headers
)
assert response.status_code == 201
assert response.json()["name"] == vector_store.name
assert response.json()["data"] == vector_store.data
return response.json()

1283
tests/data/Vector_store.json Normal file

File diff suppressed because it is too large Load diff

View file

@ -451,6 +451,23 @@ def test_basic_chat_with_two_session_ids_and_names(client, added_flow, created_a
assert name in response_json["result"]["text"]
def test_vector_store_in_process(client, added_vector_store, created_api_key):
# Run the /api/v1/process/{flow_id} endpoint
headers = {"x-api-key": created_api_key.api_key}
post_data = {"inputs": {"input": "What is Langflow?"}}
response = client.post(
f"api/v1/process/{added_vector_store.get('id')}",
headers=headers,
json=post_data,
)
assert response.status_code == 200, response.json()
# Check the response
assert "Langflow" in response.json()["result"]["output"]
# session_id should be returned
assert "session_id" in response.json()
assert response.json()["session_id"] is not None
# Test function without loop
@pytest.mark.async_test
def test_async_task_processing(client, added_flow, created_api_key):
@ -477,3 +494,37 @@ def test_async_task_processing(client, added_flow, created_api_key):
assert "result" in task_status_json, task_status_json
assert "text" in task_status_json["result"], task_status_json["result"]
assert "Gabriel" in task_status_json["result"]["text"], task_status_json["result"]
# Test function without loop
@pytest.mark.async_test
def test_async_task_processing_vector_store(
client, added_vector_store, created_api_key
):
headers = {"x-api-key": created_api_key.api_key}
post_data = {"inputs": {"input": "What is Langflow?"}}
# Run the /api/v1/process/{flow_id} endpoint with sync=False
response = client.post(
f"api/v1/process/{added_vector_store.get('id')}",
headers=headers,
json={**post_data, "sync": False},
)
assert response.status_code == 200, response.json()
assert "result" in response.json()
assert "FAILURE" not in response.json()["result"]
# Extract the task ID from the response
task_id = response.json().get("id")
assert task_id is not None
# Polling the task status using the helper function
task_status_json = poll_task_status(client, headers, task_id, max_attempts=40)
assert task_status_json is not None, "Task did not complete in time"
# Validate that the task completed successfully and the result is as expected
assert "result" in task_status_json, task_status_json
assert "output" in task_status_json["result"], task_status_json["result"]
assert "Langflow is" in task_status_json["result"]["output"], task_status_json[
"result"
]