🔧 chore(.gitignore): add .docker to the list of ignored files

🐛 fix(locustfile.py): remove unused import and fix incorrect variable name
✨ feat(locustfile.py): add support for authentication and flow creation in on_start method

.gitignore

@@ -254,3 +254,4 @@ langflow.db
 
 /tmp/*
 src/backend/langflow/frontend/
+.docker

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 <name>"
             },
             "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"]