From 44da36eb2df4c8d0b23b07e6d4062870fa7dc757 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Thu, 18 Jun 2026 19:25:01 +0000 Subject: [PATCH] feat: add load tests for conversational analytics - Added `cookieless-embed-conversational-analytics` load test for UI-based testing. - Added `conversational-analytics-api` load test for direct API testing. - Integrated new commands into the `lkr` CLI with support for targeting agents, custom prompts, and multiple questions. - Implemented cookieless embed handshake server and container for conversational analytics. Co-authored-by: bwebs <14831748+bwebs@users.noreply.github.com> --- .../embed_container.html | 160 ++++++++++++ .../embed_server.py | 162 ++++++++++++ ...locustfile_conversational_analytics_api.py | 143 +++++++++++ ...okieless_embed_conversational_analytics.py | 152 +++++++++++ lkr/main.py | 236 +++++++++++++++++- 5 files changed, 852 insertions(+), 1 deletion(-) create mode 100644 lkr/load_test/embed_cookieless_conversational_analytics/embed_container.html create mode 100644 lkr/load_test/embed_cookieless_conversational_analytics/embed_server.py create mode 100644 lkr/load_test/locustfile_conversational_analytics_api.py create mode 100644 lkr/load_test/locustfile_cookieless_embed_conversational_analytics.py diff --git a/lkr/load_test/embed_cookieless_conversational_analytics/embed_container.html b/lkr/load_test/embed_cookieless_conversational_analytics/embed_container.html new file mode 100644 index 0000000..1004d84 --- /dev/null +++ b/lkr/load_test/embed_cookieless_conversational_analytics/embed_container.html @@ -0,0 +1,160 @@ + + + + + + + Cookieless Embed Conversational Analytics + + + +

Cookieless Embed Conversational Analytics

+
+ + + + diff --git a/lkr/load_test/embed_cookieless_conversational_analytics/embed_server.py b/lkr/load_test/embed_cookieless_conversational_analytics/embed_server.py new file mode 100644 index 0000000..c59a842 --- /dev/null +++ b/lkr/load_test/embed_cookieless_conversational_analytics/embed_server.py @@ -0,0 +1,162 @@ +import os +import json +from http.server import BaseHTTPRequestHandler, HTTPServer +from pathlib import Path +import looker_sdk +from looker_sdk import models40 +from lkr.load_test.utils import ( + MAX_SESSION_LENGTH, + PERMISSIONS, + get_user_id, + format_attributes, +) +import sys + +class CookielessEmbedHandler(BaseHTTPRequestHandler): + def __init__(self, *args, debug=False, port=None, **kwargs): + self.debug = debug + self.sdk = looker_sdk.init40() + self.port = port + super().__init__(*args, **kwargs) + + def log_message(self, format, *args): + pass + + def do_GET(self): + if self.path == '/': + self.send_response(200) + self.send_header('Content-type', 'text/html') + self.end_headers() + html_path = Path(__file__).parent / "embed_container.html" + with open(html_path, "r") as f: + html_content = f.read() + + looker_host = os.environ.get("LOOKERSDK_BASE_URL", "") + agent_id = os.environ.get("AGENT_ID", "") + conversation_id = os.environ.get("CONVERSATION_ID", "") + + html_content = html_content.replace("{{LOOKER_HOST}}", looker_host) + html_content = html_content.replace("{{AGENT_ID}}", agent_id) + html_content = html_content.replace("{{CONVERSATION_ID}}", conversation_id) + html_content = html_content.replace("{{debug}}", str(self.debug).lower()) + + self.wfile.write(html_content.encode("utf-8")) + elif self.path == '/acquire-embed-session': + self.send_response(200) + self.send_header('Content-type', 'application/json') + self.send_header('Cache-Control', 'no-store') + self.end_headers() + + user_id = get_user_id() + + models_str = os.environ.get("MODELS", "") + models = models_str.split(",") if models_str else [] + group_ids_str = os.environ.get("GROUP_IDS", "") + group_ids = group_ids_str.split(",") if group_ids_str else [] + external_group_id = os.environ.get("EXTERNAL_GROUP_ID") + + attributes_str = os.environ.get("ATTRIBUTES", "[]") + attributes_list = json.loads(attributes_str) + user_attributes = format_attributes(attributes_list) + first_name = os.environ.get("FIRST_NAME", "Cookieless Embed") + + user_session = models40.EmbedCookielessSessionAcquire( + first_name=first_name, + last_name=user_id, + external_user_id=user_id, + session_length=3600, + permissions=PERMISSIONS + ["chat_with_explore"], + models=models, + group_ids=group_ids, + external_group_id=external_group_id, + user_attributes=user_attributes, + embed_domain=f"http://127.0.0.1:{self.port}" + ) + + try: + response = self.sdk.acquire_embed_cookieless_session( + body=user_session, + transport_options={'headers':{'User-Agent': self.headers.get('User-Agent')}} + ) + self.wfile.write(json.dumps({ + 'api_token': response.api_token, + 'api_token_ttl': response.api_token_ttl, + 'authentication_token': response.authentication_token, + 'authentication_token_ttl': response.authentication_token_ttl, + 'navigation_token': response.navigation_token, + 'navigation_token_ttl': response.navigation_token_ttl, + 'session_reference_token': response.session_reference_token, + 'session_reference_token_ttl': response.session_reference_token_ttl, + }).encode('utf-8')) + except Exception as e: + self.send_response(500) + self.end_headers() + self.wfile.write(json.dumps({'error': str(e)}).encode('utf-8')) + else: + self.send_response(404) + self.end_headers() + + def do_POST(self): + if self.path == '/generate-embed-tokens': + content_length = int(self.headers['Content-Length']) + post_data = self.rfile.read(content_length) + data = json.loads(post_data) + session_reference_token = data.get('session_reference_token') + api_token = data.get('api_token') + navigation_token = data.get('navigation_token') + + if not session_reference_token or not api_token or not navigation_token: + self.send_response(400) + self.send_header('Content-type', 'application/json') + self.end_headers() + self.wfile.write(json.dumps({'error': 'session_reference_token, api_token, and navigation_token are required'}).encode('utf-8')) + return + + try: + session_information = models40.EmbedCookielessSessionGenerateTokens( + session_reference_token=session_reference_token, + api_token=api_token, + navigation_token=navigation_token + ) + response = self.sdk.generate_tokens_for_cookieless_session( + body=session_information, + transport_options={'headers':{'User-Agent': self.headers.get('User-Agent')}} + ) + self.send_response(200) + self.send_header('Content-type', 'application/json') + self.end_headers() + self.wfile.write(json.dumps({ + 'api_token': response.api_token, + 'api_token_ttl': response.api_token_ttl, + 'navigation_token': response.navigation_token, + 'navigation_token_ttl': response.navigation_token_ttl, + 'session_reference_token': response.session_reference_token, + 'session_reference_token_ttl': response.session_reference_token_ttl, + }).encode('utf-8')) + except Exception as e: + print(e) + self.send_response(500) + self.end_headers() + self.wfile.write(json.dumps({'error': str(e)}).encode('utf-8')) + else: + self.send_response(404) + self.end_headers() + +def run_server(port=8080, debug=False): + def handler(*args, **kwargs): + CookielessEmbedHandler(*args, debug=debug, port=port, **kwargs) + + server_address = ('' , port) + httpd = HTTPServer(server_address, handler) + httpd.serve_forever() + +if __name__ == '__main__': + port = 8080 + debug = "--debug" in sys.argv + if len(sys.argv) > 1 and not sys.argv[1].startswith("--"): + try: + port = int(sys.argv[1]) + except ValueError: + print("Invalid port number", file=sys.stderr) + sys.exit(1) + run_server(port, debug) diff --git a/lkr/load_test/locustfile_conversational_analytics_api.py b/lkr/load_test/locustfile_conversational_analytics_api.py new file mode 100644 index 0000000..f369cf0 --- /dev/null +++ b/lkr/load_test/locustfile_conversational_analytics_api.py @@ -0,0 +1,143 @@ +import datetime +import os +import random +from typing import List, Optional + +import looker_sdk +from locust import User, between, task +from looker_sdk import models40 +from looker_sdk.sdk.api40.methods import Looker40SDK +from structlog import get_logger + +from lkr.load_test.utils import ( + MAX_SESSION_LENGTH, + PERMISSIONS, + extract_looker_user_id_from_token, + format_attributes, + get_user_id, +) + +logger = get_logger(__name__) + +class ConversationalAnalyticsApiUser(User): + abstract = True + wait_time = between(1, 5) + host = os.environ.get("LOOKERSDK_BASE_URL") + cleanup_user: bool = True + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.sdk: Looker40SDK | None = None + self.user_id = get_user_id() + self.agent_id: Optional[str] = None + self.models: List[str] = [] + self.explores: List[str] = [] + self.agent_prompt: Optional[str] = None + self.questions: List[str] = ["What are the top 5 products by sales?"] + self.continue_conversation: bool = False + self.attributes: List[str] = [] + self.group_ids: List[str] = [] + self.external_group_id: str | None = None + self.first_name: str = "Embed" + self.conversation_id: Optional[str] = None + + def _init_sdk(self): + sdk = looker_sdk.init40() + attributes = format_attributes(self.attributes) + embed_session = sdk.acquire_embed_cookieless_session( + models40.EmbedCookielessSessionAcquire( + first_name=self.first_name, + last_name=self.user_id, + external_user_id=self.user_id, + external_group_id=self.external_group_id, + session_length=MAX_SESSION_LENGTH, + permissions=PERMISSIONS + ["chat_with_explore"], + models=self.models, + user_attributes=attributes, + group_ids=self.group_ids, + ) + ) + looker_user_id = extract_looker_user_id_from_token(embed_session) + if not looker_user_id: + embed_user = sdk.user_for_credential("embed", self.user_id) + if not embed_user or not embed_user.id: + raise Exception("Failed to create embed user") + looker_user_id = int(embed_user.id) + + sdk.auth.login_user(looker_user_id) + return sdk + + def _get_or_create_agent(self, sdk: Looker40SDK): + if self.agent_id: + return self.agent_id + + # Create a new agent + sources = [] + for model in self.models: + for explore in self.explores: + sources.append(models40.Source(model=model, explore=explore)) + + if not sources: + # Fallback if no explores specified but models are + for model in self.models: + sources.append(models40.Source(model=model)) + + write_agent = models40.WriteAgent( + name=f"Load Test Agent {self.user_id}", + description="Agent created for load testing", + sources=sources, + context=models40.Context(instructions=self.agent_prompt) if self.agent_prompt else None + ) + agent = sdk.create_agent(body=write_agent) + return agent.id + + def _create_conversation(self, sdk: Looker40SDK, agent_id: str): + write_conversation = models40.WriteConversation( + name=f"Load Test Conversation {self.user_id}", + agent_id=agent_id + ) + conversation = sdk.create_conversation(body=write_conversation) + return conversation.id + + def on_start(self): + self.sdk = self._init_sdk() + self.agent_id = self._get_or_create_agent(self.sdk) + if self.continue_conversation: + self.conversation_id = self._create_conversation(self.sdk, self.agent_id) + + @task + def chat(self): + if not self.sdk or not self.agent_id: + return + + question = random.choice(self.questions) + + cid = self.conversation_id + if not cid: + cid = self._create_conversation(self.sdk, self.agent_id) + + start_time = datetime.datetime.now() + try: + self.sdk.conversational_analytics_chat( + body=models40.ConversationalAnalyticsChatRequest( + conversation_id=cid, + user_message=question + ) + ) + end_time = datetime.datetime.now() + duration = (end_time - start_time).total_seconds() + + logger.info( + "conversational_analytics_chat", + question=question, + duration=duration, + conversation_id=cid, + agent_id=self.agent_id + ) + except Exception as e: + logger.error("chat_failed", error=str(e), question=question, conversation_id=cid) + + if not self.continue_conversation: + # If not continuing, we don't save the conversation_id for the next task + # (In standard Looker, we might want to delete it too, but maybe not for load test) + pass diff --git a/lkr/load_test/locustfile_cookieless_embed_conversational_analytics.py b/lkr/load_test/locustfile_cookieless_embed_conversational_analytics.py new file mode 100644 index 0000000..b454f03 --- /dev/null +++ b/lkr/load_test/locustfile_cookieless_embed_conversational_analytics.py @@ -0,0 +1,152 @@ +from locust import User, between, task +import os +import socket +import subprocess +import sys +import time +import json +from typing import List +from selenium import webdriver +from selenium.webdriver.chrome.options import Options +from selenium.webdriver.common.by import By +from selenium.webdriver.support.ui import WebDriverWait +from selenium.webdriver.support import expected_conditions as EC +from urllib.parse import urlparse + +def get_free_port(): + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + s.close() + return port + +class CookielessEmbedConversationalAnalyticsUser(User): + abstract = True + wait_time = between(1000, 2000) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.port = get_free_port() + self.host = f"http://127.0.0.1:{self.port}" + self.debug = getattr(self, 'debug', False) + self.group_ids: List[str] = getattr(self, 'group_ids', []) + self.external_group_id: str | None = getattr(self, 'external_group_id', None) + self.agent_id: str = getattr(self, 'agent_id', "") + self.conversation_id: str = getattr(self, 'conversation_id', "") + self.models: List[str] = getattr(self, 'models', []) + self.attributes: List[str] = getattr(self, 'attributes', []) + self.first_name: str = getattr(self, 'first_name', 'Embed') + + server_path = os.path.join( + os.path.dirname(__file__), + "embed_cookieless_conversational_analytics", + "embed_server.py" + ) + + server_cmd = [sys.executable, server_path, str(self.port)] + if self.debug: + server_cmd.append("--debug") + + lEnv = os.environ.copy() + lEnv["AGENT_ID"] = self.agent_id + lEnv["CONVERSATION_ID"] = self.conversation_id + lEnv["MODELS"] = ",".join(self.models) + lEnv["GROUP_IDS"] = ",".join(self.group_ids) + lEnv["ATTRIBUTES"] = json.dumps(self.attributes) + lEnv["FIRST_NAME"] = self.first_name + if self.external_group_id: + lEnv["EXTERNAL_GROUP_ID"] = self.external_group_id + + # Fail fast in Python: set timeout for Looker SDK requests + if "LOOKERSDK_TIMEOUT" not in lEnv: + lEnv["LOOKERSDK_TIMEOUT"] = "10" + + self.server_process = subprocess.Popen( + server_cmd, env=lEnv + ) + + # Wait for the server to be ready with 10 second timeout + is_server_ready = False + for _ in range(20): + try: + with socket.create_connection(("127.0.0.1", self.port), timeout=0.5): + is_server_ready = True + break + except (socket.timeout, ConnectionRefusedError): + time.sleep(0.5) + + if not is_server_ready: + self.server_process.terminate() + self.server_process.wait() + raise Exception("Embed server failed to start") + + chrome_options = Options() + chrome_options.add_argument("--headless=new") + chrome_options.add_argument("--no-sandbox") + chrome_options.add_argument("--disable-dev-shm-usage") + chrome_options.add_argument("--disable-gpu") + chrome_options.add_argument("--user-agent=Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36") + chrome_options.add_argument("--enable-logging") + chrome_options.add_argument("--v=1") + + # Speed up page loading by not waiting for non-critical subresources (CSS, images, fonts) + chrome_options.page_load_strategy = "eager" + + # In VPCSC, block everything except required hosts to trigger immediate failure + # instead of waiting for a 60-second network timeout. + looker_url = os.environ.get("LOOKERSDK_BASE_URL", "") + looker_host = urlparse(looker_url).hostname + + rules = "MAP * ~NOTFOUND, EXCLUDE localhost, EXCLUDE 127.0.0.1" + if looker_host: + rules += f", EXCLUDE {looker_host}" + + chrome_options.add_argument(f"--host-resolver-rules={rules}") + + chrome_options.add_experimental_option( + "prefs", + { + "profile.cookie_controls_mode": 1, # 1 = Block third-party cookies + }, + ) + if self.debug: + chrome_options.set_capability("goog:loggingPrefs", {"browser": "ALL"}) + self.driver = webdriver.Chrome(options=chrome_options) + + def on_start(self): + self.driver.get(self.host) + try: + # Waiting up to 2 seconds for embed iframe to be present...") + WebDriverWait(self.driver, 2).until( + EC.presence_of_element_located((By.ID, "looker-embed")) + ) + print("Embed iframe is present. Waiting 2 seconds for handshake to complete...") + time.sleep(2) + except Exception as e: + print(f"Error waiting for iframe or handshake: {e}") + finally: + for entry in self.driver.get_log('browser'): + print(entry) + + def on_stop(self): + try: + if hasattr(self, "driver") and self.driver: + self.driver.quit() + except BaseException as e: + print(f"Notice: Exception during driver.quit(): {e}") + finally: + self.driver = None + + try: + if hasattr(self, "server_process") and self.server_process: + self.server_process.terminate() + self.server_process.wait() + except BaseException as e: + print(f"Notice: Exception during server cleanup: {e}") + finally: + self.server_process = None + + @task + def do_nothing(self): + pass diff --git a/lkr/main.py b/lkr/main.py index 63f4428..1cc94c4 100644 --- a/lkr/main.py +++ b/lkr/main.py @@ -26,6 +26,8 @@ from lkr.load_test.locustfile_qid import QueryUser from lkr.load_test.locustfile_render import RenderUser from lkr.load_test.locustfile_cookieless_embed_dashboard import CookielessEmbedDashboardUser +from lkr.load_test.locustfile_cookieless_embed_conversational_analytics import CookielessEmbedConversationalAnalyticsUser +from lkr.load_test.locustfile_conversational_analytics_api import ConversationalAnalyticsApiUser from lkr.load_test.locustfile_dashboard_queries import DashboardQueriesUser from lkr.load_test.utils import get_external_group_id, get_dashboard_load_test_system_activity_explore_url from lkr.utils.validate_api import validate_api_credentials @@ -52,6 +54,8 @@ class LoadTestType(str, Enum): render = "render" cookieless_embed = "cookieless-embed" cookieless_embed_dashboard = "cookieless-embed-dashboard" + cookieless_embed_conversational_analytics = "cookieless-embed-conversational-analytics" + conversational_analytics_api = "conversational-analytics-api" class DebugType(str, Enum): @@ -153,6 +157,8 @@ def check_settings( "embed-observability", "cookieless-embed", "cookieless-embed-dashboard", + "cookieless-embed-conversational-analytics", + "conversational-analytics-api", "dashboard-queries", ]: sdk = looker_sdk.init40() @@ -175,7 +181,7 @@ def check_settings( ) raise typer.Exit(1) - if ctx.invoked_subcommand in ["query", "render", "cookieless-embed", "cookieless-embed-dashboard", "dashboard-queries"]: + if ctx.invoked_subcommand in ["query", "render", "cookieless-embed", "cookieless-embed-dashboard", "cookieless-embed-conversational-analytics", "conversational-analytics-api", "dashboard-queries"]: # check for embed cookieless v2 if not setting.embed_cookieless_v2: typer.echo( @@ -325,6 +331,234 @@ def quit_runner(): runner.spawning_greenlet.spawn_later(run_time * 60, quit_runner) runner.greenlet.join() + +@group.command(name="cookieless-embed-conversational-analytics") +def load_test_cookieless_embed_conversational_analytics( + agent_id: str = typer.Option( + help="Agent ID to run the test on", + default="", + ), + conversation_id: str = typer.Option( + help="Conversation ID to run the test on", + default="", + ), + model: list[str] = typer.Option( + help="Model to run the test on. Specify multiple models as --model model1 --model model2", + default=..., + ), + attribute: Annotated[ + List[str] | None, + typer.Option( + help="Looker attributes to run the test on. Specify them as attribute:value like --attribute store:value." + ), + ] = None, + group: Annotated[ + List[str], + typer.Option( + help="Looker group IDs to add to the user." + ), ] = [], + external_group_id: Annotated[ + str | None, + typer.Option( + help="External group ID to add to the user. Will be prefixed with embed unless overridden with --external-group-id-prefix" + ), + ] = None, + external_group_id_prefix: Annotated[ + str | None, + typer.Option( + help="Prefix to add to the group IDs. Defaults to `embed`. To remove the prefix, pass in an empty string" + ), + ] = "embed", + users: Annotated[ + int, typer.Option(help="Number of users to run the test with", min=1, max=1000) + ] = 25, + spawn_rate: Annotated[ + float, + typer.Option(help="Number of users to spawn per second", min=0, max=100), + ] = 1, + run_time: Annotated[ + int, typer.Option(help="How many minutes to run the load test for", min=1) + ] = 5, + stop_timeout: Annotated[ + int, + typer.Option( + help="How many seconds to wait for the load test to stop", + ), + ] = 15, + debug: Annotated[ + bool, + typer.Option( + "--debug", + help="Enable debug mode", + ), + ] = False, + first_name: Annotated[ + str, + typer.Option( + help="First name of the embed user", + ), + ] = "Embed", +): + """ + Run a load test on Conversational Analytics using Cookieless Embed V2. + """ + from locust import events + from locust.env import Environment + + typer.echo( + f"Running load test with {users} users, {spawn_rate} spawn rate, and {run_time} minutes" + ) + + class CookielessEmbedConversationalAnalyticsUserClass(CookielessEmbedConversationalAnalyticsUser): + def __init__(self, *args, **kwargs): + self.debug = debug + self.agent_id = agent_id + self.conversation_id = conversation_id + self.models = model + self.attributes = attribute or [] + self.group_ids = group or [] + self.external_group_id = get_external_group_id( + external_group_id, external_group_id_prefix + ) + self.first_name = first_name + super().__init__(*args, **kwargs) + + env = Environment( + user_classes=[CookielessEmbedConversationalAnalyticsUserClass], events=events, stop_timeout=stop_timeout + ) + runner = env.create_local_runner() + + runner.start(user_count=users, spawn_rate=spawn_rate) + + def quit_runner(): + runner.stop() + if runner.greenlet: + runner.greenlet.kill(block=False) + typer.Exit(1) + + if runner.spawning_greenlet: + runner.spawning_greenlet.spawn_later(run_time * 60, quit_runner) + runner.greenlet.join() + + +@group.command(name="conversational-analytics-api") +def load_test_conversational_analytics_api( + agent_id: str = typer.Option( + help="Target with an agent id", + default="", + ), + agent_prompt: str = typer.Option( + help="Give it an agent prompt and it creates a new one", + default="", + ), + model: list[str] = typer.Option( + help="Model to run the test on. Specify multiple models as --model model1 --model model2", + default=..., + ), + explore: list[str] = typer.Option( + help="Explores to include if agent doesn't exist yet. Specify multiple as --explore exp1 --explore exp2", + default=[], + ), + question: list[str] = typer.Option( + help="Questions to ask the agent. Specify multiple as --question 'q1' --question 'q2'", + default=["What are the top 5 products by sales?"], + ), + continue_conversation: Annotated[ + bool, + typer.Option( + help="Leaves conversation open and sends more chats", + ), + ] = False, + attribute: Annotated[ + List[str] | None, + typer.Option( + help="Looker attributes to run the test on. Specify them as attribute:value like --attribute store:value." + ), + ] = None, + group: Annotated[ + List[str], + typer.Option( + help="Looker group IDs to add to the user." + ), ] = [], + external_group_id: Annotated[ + str | None, + typer.Option( + help="External group ID to add to the user. Will be prefixed with embed unless overridden with --external-group-id-prefix" + ), + ] = None, + external_group_id_prefix: Annotated[ + str | None, + typer.Option( + help="Prefix to add to the group IDs. Defaults to `embed`. To remove the prefix, pass in an empty string" + ), + ] = "embed", + users: Annotated[ + int, typer.Option(help="Number of users to run the test with", min=1, max=1000) + ] = 25, + spawn_rate: Annotated[ + float, + typer.Option(help="Number of users to spawn per second", min=0, max=100), + ] = 1, + run_time: Annotated[ + int, typer.Option(help="How many minutes to run the load test for", min=1) + ] = 5, + stop_timeout: Annotated[ + int, + typer.Option( + help="How many seconds to wait for the load test to stop", + ), + ] = 15, + first_name: Annotated[ + str, + typer.Option( + help="First name of the embed user", + ), + ] = "Embed", +): + """ + Run a load test on Conversational Analytics using the Looker API. + """ + from locust import events + from locust.env import Environment + + typer.echo( + f"Running API load test with {users} users, {spawn_rate} spawn rate, and {run_time} minutes" + ) + + class ConversationalAnalyticsApiUserClass(ConversationalAnalyticsApiUser): + def __init__(self, *args, **kwargs): + self.agent_id = agent_id + self.agent_prompt = agent_prompt + self.models = model + self.explores = explore + self.questions = question + self.continue_conversation = continue_conversation + self.attributes = attribute or [] + self.group_ids = group or [] + self.external_group_id = get_external_group_id( + external_group_id, external_group_id_prefix + ) + self.first_name = first_name + super().__init__(*args, **kwargs) + + env = Environment( + user_classes=[ConversationalAnalyticsApiUserClass], events=events, stop_timeout=stop_timeout + ) + runner = env.create_local_runner() + + runner.start(user_count=users, spawn_rate=spawn_rate) + + def quit_runner(): + runner.stop() + if runner.greenlet: + runner.greenlet.kill(block=False) + typer.Exit(1) + + if runner.spawning_greenlet: + runner.spawning_greenlet.spawn_later(run_time * 60, quit_runner) + runner.greenlet.join() + + @group.command(name="dashboard") def load_test( dashboard: str = typer.Option(