diff --git a/.gitignore b/.gitignore index 156f44394..9b8ec1cc1 100644 --- a/.gitignore +++ b/.gitignore @@ -254,3 +254,4 @@ langflow.db /tmp/* src/backend/langflow/frontend/ +.docker \ No newline at end of file diff --git a/tests/locust/locustfile.py b/tests/locust/locustfile.py index 765637a61..259c580c8 100644 --- a/tests/locust/locustfile.py +++ b/tests/locust/locustfile.py @@ -1,21 +1,20 @@ -from locust import FastHttpUser, task, between +from locust import FastHttpUser, task, between, events import random import time +import orjson from rich import print import os +import httpx +from pathlib import Path class NameTest(FastHttpUser): - host = "http://localhost:7860/api/v1" # make sure the port number is correct wait_time = between(1, 5) with open("names.txt", "r") as file: names = [line.strip() for line in file.readlines()] - headers = { - # api-key - "x-api-key": "lf-a1EsRC75ybWiKRCRLyhX1R9rPlqQKYlnRQoZWysg6NM", - } + headers = {} def poll_task(self, task_id, sleep_time=1): while True: @@ -26,53 +25,109 @@ class NameTest(FastHttpUser): headers=self.headers, ) as response: status = response.js.get("status") + print(f"Poll Response: {response.js}") if status == "SUCCESS": return response.js.get("result") elif status in ["FAILURE", "REVOKED"]: raise ValueError(f"Task failed with status: {status}") time.sleep(sleep_time) + def process(self, name, flow_id, payload): + task_id = None + print(f"Processing {payload}") + with self.rest( + "POST", + f"/process/{flow_id}", + json=payload, + name="process", + headers=self.headers, + ) as response: + print(response.js) + if response.status_code != 200: + response.failure("Process call failed") + raise ValueError("Process call failed") + task_id = response.js.get("id") + session_id = response.js.get("session_id") + assert task_id, "Inner Task ID not found" + + assert task_id, "Task ID not found" + result = self.poll_task(task_id) + print(f"Result for {name}: {result}") + + return result, session_id + @task def send_name_and_check(self): name = random.choice(self.names) - flow_id = os.getenv("FLOW_ID", "88989f6d-8da4-4d6f-8205-480693c36200") - session_id = f"{name}-{time.time()}" - - def process(flow_id, payload): - task_id = None - print(f"Processing {payload}") - with self.rest( - "POST", - f"/process/{flow_id}", - json=payload, - name="process", - headers=self.headers, - ) as response: - if response.status_code != 200: - response.failure("Process call failed") - raise ValueError("Process call failed") - print(response.js) - task_id = response.js.get("id") - assert task_id, "Inner Task ID not found" - - assert task_id, "Task ID not found" - result, session_id = self.poll_task(task_id) - print(f"Result for {name}: {result}") - - return result, session_id payload1 = { "inputs": {"text": f"Hello, My name is {name}"}, - "session_id": session_id, + "sync": False, } - result1, session_id = process(flow_id, payload1) + result1, session_id = self.process(name, self.flow_id, payload1) payload2 = { "inputs": { "text": "What is my name? Please, answer like this: Your name is " }, "session_id": session_id, + "sync": False, } - result2, session_id = process(flow_id, payload2) + result2, session_id = self.process(name, self.flow_id, payload2) - assert f"Your name is {name}" in result2, "Name not found in response" + assert f"Your name is {name}" in str(result2), "Name not found in response" + + def on_start(self): + print("Starting") + login_data = {"username": "superuser", "password": "superuser"} + response = httpx.post(f"{self.host}/login", data=login_data) + print(response.json()) + + tokens = response.json() + print(tokens) + a_token = tokens["access_token"] + logged_in_headers = {"Authorization": f"Bearer {a_token}"} + print("Logged in") + with open( + Path(__file__).parent.parent + / "data" + / "BasicChatwithPromptandHistory.json", + "r", + ) as f: + json_flow = f.read() + flow = orjson.loads(json_flow) + data = flow["data"] + # Create test data + flow = {"name": "Flow 1", "description": "description", "data": data} + print("Creating flow") + # Make request to endpoint + response = httpx.post( + f"{self.host}/flows/", + json=flow, + headers=logged_in_headers, + ) + self.flow_id = response.json()["id"] + print(f"Flow ID: {self.flow_id}") + + # read all users + response = httpx.get( + f"{self.host}/users/", + headers=logged_in_headers, + ) + print(response.json()) + user_id = next( + ( + user["id"] + for user in response.json()["users"] + if user["username"] == "superuser" + ), + None, + ) + # Create api key + response = httpx.post( + f"{self.host}/api_key/", + json={"user_id": user_id}, + headers=logged_in_headers, + ) + print(response.json()) + self.headers["x-api-key"] = response.json()["api_key"]