diff --git a/scripts/stress-test/README.md b/scripts/stress-test/README.md index 15f21cd5329..f3a5827ed11 100644 --- a/scripts/stress-test/README.md +++ b/scripts/stress-test/README.md @@ -84,10 +84,10 @@ The stress test tests a single endpoint with comprehensive SSE metrics tracking: ## Prerequisites -1. **Dependencies are automatically installed** when running setup: +1. **Dependencies**: - - Locust (load testing framework) - - sseclient-py (SSE client library) + - Locust runs through `uvx --from locust`, outside the API project environment. + - `sseclient-py` is included in the API project dependencies. 1. **Complete Dify setup**: @@ -96,6 +96,25 @@ The stress test tests a single endpoint with comprehensive SSE metrics tracking: python scripts/stress-test/setup_all.py ``` + For a brand-new Dify instance, the setup script creates the first admin account. Override the defaults if needed: + + ```bash + STRESS_TEST_ADMIN_EMAIL='your-admin@example.com' \ + STRESS_TEST_ADMIN_USERNAME='dify' \ + STRESS_TEST_ADMIN_PASSWORD='your-password' \ + python scripts/stress-test/setup_all.py + ``` + + For an already-initialized Dify instance with an admin account, provide the existing admin login: + + ```bash + STRESS_TEST_ADMIN_EMAIL='your-admin@example.com' \ + STRESS_TEST_ADMIN_PASSWORD='your-password' \ + python scripts/stress-test/setup_all.py + ``` + + `STRESS_TEST_ADMIN_USERNAME` is only used in the brand-new instance case, when `/console/api/setup` creates the first admin account. + 1. **Ensure services are running**: **IMPORTANT**: For accurate stress testing, run the API server with Gunicorn in production mode: @@ -141,11 +160,11 @@ The stress test tests a single endpoint with comprehensive SSE metrics tracking: # Run with default configuration (headless mode) ./scripts/stress-test/run_locust_stress_test.sh -# Or run directly with uv -uv run --project api python -m locust -f scripts/stress-test/sse_benchmark.py --host http://localhost:5001 +# Or run directly with uvx +uvx --from locust locust -f scripts/stress-test/sse_benchmark.py --host http://localhost:5001 # Run with Web UI (access at http://localhost:8089) -uv run --project api python -m locust -f scripts/stress-test/sse_benchmark.py --host http://localhost:5001 --web-port 8089 +uvx --from locust locust -f scripts/stress-test/sse_benchmark.py --host http://localhost:5001 --web-port 8089 ``` The script will: @@ -182,12 +201,13 @@ self.questions = [ ### Report Structure -After running the stress test, you'll find these files in the `reports/` directory: +After running the stress test, you'll find one directory per run under `reports/`: -- `locust_summary_YYYYMMDD_HHMMSS.txt` - Complete console output with metrics -- `locust_report_YYYYMMDD_HHMMSS.html` - Interactive HTML report with charts -- `locust_YYYYMMDD_HHMMSS_stats.csv` - CSV with detailed statistics -- `locust_YYYYMMDD_HHMMSS_stats_history.csv` - Time-series data +- `YYYYMMDD_HHMMSS/locust_summary.txt` - Complete console output with metrics +- `YYYYMMDD_HHMMSS/locust_report.html` - Interactive HTML report with charts +- `YYYYMMDD_HHMMSS/locust_stats.csv` - CSV with detailed statistics +- `YYYYMMDD_HHMMSS/locust_stats_history.csv` - Time-series data +- `YYYYMMDD_HHMMSS/sse_metrics_YYYYMMDD_HHMMSS.json` - Custom SSE metrics ### Key Metrics @@ -399,8 +419,8 @@ docker compose -f docker/docker-compose.middleware.yaml up -d db 1. **"ModuleNotFoundError: No module named 'locust'"**: ```bash - # Dependencies are installed automatically, but if needed: - uv --project api add --dev locust sseclient-py + # Locust is intentionally run outside the api project environment: + uvx --from locust locust --version ``` 1. **"API key configuration not found"**: @@ -453,15 +473,15 @@ Run Locust directly with custom options: ```bash # With specific user count and spawn rate -uv run --project api python -m locust -f scripts/stress-test/sse_benchmark.py \ +uvx --from locust locust -f scripts/stress-test/sse_benchmark.py \ --host http://localhost:5001 --users 50 --spawn-rate 5 # Generate CSV reports -uv run --project api python -m locust -f scripts/stress-test/sse_benchmark.py \ +uvx --from locust locust -f scripts/stress-test/sse_benchmark.py \ --host http://localhost:5001 --csv reports/results # Run for specific duration -uv run --project api python -m locust -f scripts/stress-test/sse_benchmark.py \ +uvx --from locust locust -f scripts/stress-test/sse_benchmark.py \ --host http://localhost:5001 --run-time 5m --headless ``` @@ -469,7 +489,7 @@ uv run --project api python -m locust -f scripts/stress-test/sse_benchmark.py \ ```bash # Compare multiple stress test runs -ls -la reports/stress_test_*.txt | tail -5 +ls -la scripts/stress-test/reports/*/locust_summary.txt | tail -5 ``` ## Interpreting Performance Issues diff --git a/scripts/stress-test/common/config_helper.py b/scripts/stress-test/common/config_helper.py index fffb5e00d80..4038259c736 100644 --- a/scripts/stress-test/common/config_helper.py +++ b/scripts/stress-test/common/config_helper.py @@ -8,9 +8,9 @@ from typing import NotRequired, TypedDict class AdminConfig(TypedDict): """Configuration for admin section.""" + email: str username: str password: str - base_url: str class AuthConfig(TypedDict): @@ -18,6 +18,7 @@ class AuthConfig(TypedDict): access_token: str refresh_token: NotRequired[str] + csrf_token: NotRequired[str] expires_at: NotRequired[int] @@ -253,18 +254,25 @@ class ConfigHelper: Returns: Access token string or None if not found """ - auth = self.get_state_section[AuthConfig]("auth") + auth = self.get_state_section("auth") if auth: return auth.get("access_token") return None + def get_csrf_token(self) -> str | None: + """Get the CSRF token from auth section.""" + auth = self.get_state_section("auth") + if auth: + return auth.get("csrf_token") + return None + def get_app_id(self) -> str | None: """Get the app ID from app section. Returns: App ID string or None if not found """ - app = self.get_state_section[AppConfig]("app") + app = self.get_state_section("app") if app: return app.get("app_id") return None @@ -275,11 +283,31 @@ class ConfigHelper: Returns: API key token string or None if not found """ - api_key = self.get_state_section[ApiKeyConfig]("api_key") + api_key = self.get_state_section("api_key") if api_key: return api_key.get("token") return None + def console_auth_headers(self) -> dict[str, str]: + access_token = self.get_token() + csrf_token = self.get_csrf_token() + headers: dict[str, str] = {} + if access_token: + headers["authorization"] = f"Bearer {access_token}" + if csrf_token: + headers["X-CSRF-Token"] = csrf_token + return headers + + def console_auth_cookies(self) -> dict[str, str]: + access_token = self.get_token() + csrf_token = self.get_csrf_token() + cookies = {"locale": "en-US"} + if access_token: + cookies["access_token"] = access_token + if csrf_token: + cookies["csrf_token"] = csrf_token + return cookies + # Create a default instance for convenience config_helper = ConfigHelper() diff --git a/scripts/stress-test/run_locust_stress_test.sh b/scripts/stress-test/run_locust_stress_test.sh index 665cb68754c..a861ada701c 100755 --- a/scripts/stress-test/run_locust_stress_test.sh +++ b/scripts/stress-test/run_locust_stress_test.sh @@ -21,10 +21,12 @@ NC='\033[0m' # No Color # Configuration TIMESTAMP=$(date +"%Y%m%d_%H%M%S") -REPORT_DIR="${STRESS_TEST_DIR}/reports" -CSV_PREFIX="${REPORT_DIR}/locust_${TIMESTAMP}" -HTML_REPORT="${REPORT_DIR}/locust_report_${TIMESTAMP}.html" -SUMMARY_REPORT="${REPORT_DIR}/locust_summary_${TIMESTAMP}.txt" +START_EPOCH=$(date +%s) +REPORT_ROOT="${STRESS_TEST_DIR}/reports" +REPORT_DIR="${REPORT_ROOT}/${TIMESTAMP}" +CSV_PREFIX="${REPORT_DIR}/locust" +HTML_REPORT="${REPORT_DIR}/locust_report.html" +SUMMARY_REPORT="${REPORT_DIR}/locust_summary.txt" # Create reports directory if it doesn't exist mkdir -p "${REPORT_DIR}" @@ -111,6 +113,7 @@ echo # Use SSE stress test script LOCUST_SCRIPT="${STRESS_TEST_DIR}/sse_benchmark.py" +LOCUST_RUN=(uvx --from locust locust) # Prepare Locust command if [ "$choice" = "2" ]; then @@ -119,7 +122,7 @@ if [ "$choice" = "2" ]; then echo # Run with web UI - uv --project api run locust \ + "${LOCUST_RUN[@]}" \ -f ${LOCUST_SCRIPT} \ --host http://localhost:5001 \ --web-port 8089 @@ -128,7 +131,7 @@ else echo # Run in headless mode with CSV output - uv --project api run locust \ + "${LOCUST_RUN[@]}" \ -f ${LOCUST_SCRIPT} \ --host http://localhost:5001 \ --users $USERS \ @@ -139,6 +142,22 @@ else --csv=$CSV_PREFIX \ --html=$HTML_REPORT \ 2>&1 | tee $SUMMARY_REPORT + SSE_METRICS_REPORT=$(python3 - <= ${START_EPOCH} +] +if reports: + source = max(reports, key=lambda p: p.stat().st_mtime) + target = Path("${REPORT_DIR}") / source.name + if source != target: + shutil.move(str(source), str(target)) + print(target) +EOF +) echo echo -e "${GREEN}═══════════════════════════════════════════════════════════════${NC}" @@ -150,6 +169,11 @@ else echo -e " ${YELLOW}HTML Report:${NC} $HTML_REPORT" echo -e " ${YELLOW}CSV Stats:${NC} ${CSV_PREFIX}_stats.csv" echo -e " ${YELLOW}CSV History:${NC} ${CSV_PREFIX}_stats_history.csv" + if [ -n "$SSE_METRICS_REPORT" ]; then + echo -e " ${YELLOW}SSE Metrics:${NC} $SSE_METRICS_REPORT" + else + echo -e " ${YELLOW}SSE Metrics:${NC} not found" + fi echo echo -e "${CYAN}View HTML report:${NC}" echo " open $HTML_REPORT # macOS" @@ -184,14 +208,22 @@ try: print(f" RPS: {row.get('Requests/s', 'N/A')}") break - # Show SSE-specific metrics print() print("SSE Streaming Metrics:") - for row in rows: - if 'Time to First Event' in row.get('Name', ''): - print(f" Time to First Event: {row.get('Median Response Time', 'N/A')} ms (median)") - elif 'Stream Duration' in row.get('Name', ''): - print(f" Stream Duration: {row.get('Median Response Time', 'N/A')} ms (median)") + import json + sse_metrics_report = "${SSE_METRICS_REPORT}" + if sse_metrics_report: + with open(sse_metrics_report, 'r') as metrics_file: + metrics = json.load(metrics_file).get("metrics", {}) + print(f" Total Connections: {metrics.get('total_connections', 'N/A')}") + print(f" Total Events: {metrics.get('total_events', 'N/A')}") + print(f" Connection Rate: {metrics.get('overall_conn_rate', 0):.2f} conn/s") + print(f" Event Throughput: {metrics.get('overall_event_rate', 0):.2f} events/s") + print(f" TTFE: {metrics.get('ttfe_p50', 0):.1f} ms p50 / {metrics.get('ttfe_p95', 0):.1f} ms p95") + print(f" Stream Duration: {metrics.get('stream_duration_p50', 0):.1f} ms p50 / {metrics.get('stream_duration_p95', 0):.1f} ms p95") + print(f" Inter-event Latency: {metrics.get('inter_event_latency_p50', 0):.1f} ms p50 / {metrics.get('inter_event_latency_p95', 0):.1f} ms p95") + else: + print(" No SSE metrics JSON report found for this run") except Exception as e: print(f"Could not parse metrics: {e}") @@ -199,4 +231,4 @@ EOF fi echo -e "${CYAN}═══════════════════════════════════════════════════════════════${NC}" -fi \ No newline at end of file +fi diff --git a/scripts/stress-test/setup/configure_openai_plugin.py b/scripts/stress-test/setup/configure_openai_plugin.py index 110417fa737..38a79ac7cd9 100755 --- a/scripts/stress-test/setup/configure_openai_plugin.py +++ b/scripts/stress-test/setup/configure_openai_plugin.py @@ -9,7 +9,7 @@ import httpx from common import Logger, config_helper -def configure_openai_plugin() -> None: +def configure_openai_plugin() -> bool: """Configure OpenAI plugin with mock server credentials.""" log = Logger("ConfigPlugin") @@ -20,7 +20,7 @@ def configure_openai_plugin() -> None: if not access_token: log.error("No access token found in config") log.info("Please run login_admin.py first to get access token") - return + return False log.step("Configuring OpenAI plugin with mock server...") @@ -50,14 +50,14 @@ def configure_openai_plugin() -> None: "Sec-Fetch-Mode": "cors", "Sec-Fetch-Site": "same-site", "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/139.0.0.0 Safari/537.36", - "authorization": f"Bearer {access_token}", + **config_helper.console_auth_headers(), "content-type": "application/json", "sec-ch-ua": '"Not;A=Brand";v="99", "Google Chrome";v="139", "Chromium";v="139"', "sec-ch-ua-mobile": "?0", "sec-ch-ua-platform": '"macOS"', } - cookies = {"locale": "en-US"} + cookies = config_helper.console_auth_cookies() try: # Make the configuration request @@ -73,25 +73,32 @@ def configure_openai_plugin() -> None: log.success("OpenAI plugin configured successfully!") log.key_value("API Base", config_payload["credentials"]["openai_api_base"]) log.key_value("API Key", config_payload["credentials"]["openai_api_key"]) + return True elif response.status_code == 201: log.success("OpenAI plugin credentials created successfully!") log.key_value("API Base", config_payload["credentials"]["openai_api_base"]) log.key_value("API Key", config_payload["credentials"]["openai_api_key"]) + return True elif response.status_code == 401: log.error("Configuration failed: Unauthorized") log.info("Token may have expired. Please run login_admin.py again") + return False else: log.error(f"Configuration failed with status code: {response.status_code}") log.debug(f"Response: {response.text}") + return False except httpx.ConnectError: log.error("Could not connect to Dify API at http://localhost:5001") log.info("Make sure the API server is running with: ./dev/start-api") + return False except Exception as e: log.error(f"An error occurred: {e}") + return False if __name__ == "__main__": - configure_openai_plugin() + if not configure_openai_plugin(): + sys.exit(1) diff --git a/scripts/stress-test/setup/create_api_key.py b/scripts/stress-test/setup/create_api_key.py index cd04fe57eb6..1bc42062c02 100755 --- a/scripts/stress-test/setup/create_api_key.py +++ b/scripts/stress-test/setup/create_api_key.py @@ -11,7 +11,7 @@ import httpx from common import Logger, config_helper -def create_api_key() -> None: +def create_api_key() -> bool: """Create API key for the imported app.""" log = Logger("CreateAPIKey") @@ -21,14 +21,14 @@ def create_api_key() -> None: access_token = config_helper.get_token() if not access_token: log.error("No access token found in config") - return + return False # Read app_id from config app_id = config_helper.get_app_id() if not app_id: log.error("No app_id found in config") log.info("Please run import_workflow_app.py first to import the app") - return + return False log.step(f"Creating API key for app: {app_id}") @@ -50,14 +50,14 @@ def create_api_key() -> None: "Sec-Fetch-Mode": "cors", "Sec-Fetch-Site": "same-site", "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/139.0.0.0 Safari/537.36", - "authorization": f"Bearer {access_token}", + **config_helper.console_auth_headers(), "content-type": "application/json", "sec-ch-ua": '"Not;A=Brand";v="99", "Google Chrome";v="139", "Chromium";v="139"', "sec-ch-ua-mobile": "?0", "sec-ch-ua-platform": '"macOS"', } - cookies = {"locale": "en-US"} + cookies = config_helper.console_auth_cookies() try: # Make the API key creation request @@ -91,23 +91,31 @@ def create_api_key() -> None: if config_helper.write_config("api_key_config", api_key_config): log.info(f"API key saved to: {config_helper.get_config_path('benchmark_state')}") + return True + return False else: log.error("No API token received") log.debug(f"Response: {json.dumps(response_data, indent=2)}") + return False elif response.status_code == 401: log.error("API key creation failed: Unauthorized") log.info("Token may have expired. Please run login_admin.py again") + return False else: log.error(f"API key creation failed with status code: {response.status_code}") log.debug(f"Response: {response.text}") + return False except httpx.ConnectError: log.error("Could not connect to Dify API at http://localhost:5001") log.info("Make sure the API server is running with: ./dev/start-api") + return False except Exception as e: log.error(f"An error occurred: {e}") + return False if __name__ == "__main__": - create_api_key() + if not create_api_key(): + sys.exit(1) diff --git a/scripts/stress-test/setup/import_workflow_app.py b/scripts/stress-test/setup/import_workflow_app.py index 41a76bd29be..2da40f806ac 100755 --- a/scripts/stress-test/setup/import_workflow_app.py +++ b/scripts/stress-test/setup/import_workflow_app.py @@ -11,7 +11,11 @@ import httpx from common import Logger, config_helper # type: ignore[import] -def import_workflow_app() -> None: +def is_successful_import_response(response_data: dict[str, object]) -> bool: + return response_data.get("status") != "failed" and bool(response_data.get("app_id")) + + +def import_workflow_app() -> bool: """Import workflow app from DSL file and save app_id.""" log = Logger("ImportApp") @@ -22,14 +26,14 @@ def import_workflow_app() -> None: if not access_token: log.error("No access token found in config") log.info("Please run login_admin.py first to get access token") - return + return False # Read workflow DSL file dsl_path = Path(__file__).parent / "dsl" / "workflow_llm.yml" if not dsl_path.exists(): log.error(f"DSL file not found: {dsl_path}") - return + return False with open(dsl_path) as f: yaml_content = f.read() @@ -57,14 +61,14 @@ def import_workflow_app() -> None: "Sec-Fetch-Mode": "cors", "Sec-Fetch-Site": "same-site", "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/139.0.0.0 Safari/537.36", - "authorization": f"Bearer {access_token}", + **config_helper.console_auth_headers(), "content-type": "application/json", "sec-ch-ua": '"Not;A=Brand";v="99", "Google Chrome";v="139", "Chromium";v="139"', "sec-ch-ua-mobile": "?0", "sec-ch-ua-platform": '"macOS"', } - cookies = {"locale": "en-US"} + cookies = config_helper.console_auth_cookies() try: # Make the import request @@ -79,50 +83,56 @@ def import_workflow_app() -> None: if response.status_code == 200: response_data = response.json() - # Check import status - if response_data.get("status") == "completed": + if is_successful_import_response(response_data): app_id = response_data.get("app_id") - - if app_id: - log.success("Workflow app imported successfully!") - log.key_value("App ID", app_id) - log.key_value("App Mode", response_data.get("app_mode")) - log.key_value("DSL Version", response_data.get("imported_dsl_version")) - - # Save app_id to config - app_config = { - "app_id": app_id, - "app_mode": response_data.get("app_mode"), - "app_name": "workflow_llm", - "dsl_version": response_data.get("imported_dsl_version"), - } - - if config_helper.write_config("app_config", app_config): - log.info(f"App config saved to: {config_helper.get_config_path('benchmark_state')}") - else: - log.error("Import completed but no app_id received") + if response_data.get("status") != "completed": + log.warning(f"Import status: {response_data.get('status')}") log.debug(f"Response: {json.dumps(response_data, indent=2)}") + log.success("Workflow app imported successfully!") + log.key_value("App ID", app_id) + log.key_value("App Mode", response_data.get("app_mode")) + log.key_value("DSL Version", response_data.get("imported_dsl_version")) + + # Save app_id to config + app_config = { + "app_id": app_id, + "app_mode": response_data.get("app_mode"), + "app_name": "workflow_llm", + "dsl_version": response_data.get("imported_dsl_version"), + } + + if config_helper.write_config("app_config", app_config): + log.info(f"App config saved to: {config_helper.get_config_path('benchmark_state')}") + return True + return False elif response_data.get("status") == "failed": log.error("Import failed") log.error(f"Error: {response_data.get('error')}") + return False else: - log.warning(f"Import status: {response_data.get('status')}") + log.error("Import response did not include app_id") log.debug(f"Response: {json.dumps(response_data, indent=2)}") + return False elif response.status_code == 401: log.error("Import failed: Unauthorized") log.info("Token may have expired. Please run login_admin.py again") + return False else: log.error(f"Import failed with status code: {response.status_code}") log.debug(f"Response: {response.text}") + return False except httpx.ConnectError: log.error("Could not connect to Dify API at http://localhost:5001") log.info("Make sure the API server is running with: ./dev/start-api") + return False except Exception as e: log.error(f"An error occurred: {e}") + return False if __name__ == "__main__": - import_workflow_app() + if not import_workflow_app(): + sys.exit(1) diff --git a/scripts/stress-test/setup/install_openai_plugin.py b/scripts/stress-test/setup/install_openai_plugin.py index 055e5661f81..8e8d5a634de 100755 --- a/scripts/stress-test/setup/install_openai_plugin.py +++ b/scripts/stress-test/setup/install_openai_plugin.py @@ -11,7 +11,11 @@ import httpx from common import Logger, config_helper -def install_openai_plugin() -> None: +def is_non_blocking_install_response(response_data: dict[str, object]) -> bool: + return not response_data.get("code") + + +def install_openai_plugin() -> bool: """Install OpenAI plugin using saved access token.""" log = Logger("InstallPlugin") @@ -22,7 +26,7 @@ def install_openai_plugin() -> None: if not access_token: log.error("No access token found in config") log.info("Please run login_admin.py first to get access token") - return + return False log.step("Installing OpenAI plugin...") @@ -50,14 +54,14 @@ def install_openai_plugin() -> None: "Sec-Fetch-Mode": "cors", "Sec-Fetch-Site": "same-site", "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/139.0.0.0 Safari/537.36", - "authorization": f"Bearer {access_token}", + **config_helper.console_auth_headers(), "content-type": "application/json", "sec-ch-ua": '"Not;A=Brand";v="99", "Google Chrome";v="139", "Chromium";v="139"', "sec-ch-ua-mobile": "?0", "sec-ch-ua-platform": '"macOS"', } - cookies = {"locale": "en-US"} + cookies = config_helper.console_auth_cookies() try: # Make the installation request @@ -74,8 +78,13 @@ def install_openai_plugin() -> None: task_id = response_data.get("task_id") if not task_id: + if is_non_blocking_install_response(response_data): + log.warning("No installation task returned; plugin may already be installed") + log.debug(f"Response: {response.text}") + return True log.error("No task ID received from installation request") - return + log.debug(f"Response: {response.text}") + return False log.progress(f"Installation task created: {task_id}") log.info("Polling for task completion...") @@ -103,7 +112,7 @@ def install_openai_plugin() -> None: success=False, message=f"Failed to get task status: {task_response.status_code}", ) - return + return False task_data = task_response.json() task_info = task_data.get("task", {}) @@ -119,7 +128,7 @@ def install_openai_plugin() -> None: plugin_info = plugins[0] log.key_value("Plugin ID", plugin_info.get("plugin_id")) log.key_value("Message", plugin_info.get("message")) - break + return True elif status == "failed": log.spinner_stop(success=False, message="Installation failed") @@ -128,30 +137,37 @@ def install_openai_plugin() -> None: if plugins: for plugin in plugins: log.list_item(f"{plugin.get('plugin_id')}: {plugin.get('message')}") - break + return False # Continue polling if status is "pending" or other else: log.spinner_stop(success=False, message="Installation timed out") log.error("Installation timed out after 60 seconds") + return False elif response.status_code == 401: log.error("Installation failed: Unauthorized") log.info("Token may have expired. Please run login_admin.py again") + return False elif response.status_code == 409: log.warning("Plugin may already be installed") log.debug(f"Response: {response.text}") + return True else: log.error(f"Installation failed with status code: {response.status_code}") log.debug(f"Response: {response.text}") + return False except httpx.ConnectError: log.error("Could not connect to Dify API at http://localhost:5001") log.info("Make sure the API server is running with: ./dev/start-api") + return False except Exception as e: log.error(f"An error occurred: {e}") + return False if __name__ == "__main__": - install_openai_plugin() + if not install_openai_plugin(): + sys.exit(1) diff --git a/scripts/stress-test/setup/login_admin.py b/scripts/stress-test/setup/login_admin.py index 572b8fb6500..751d850a45e 100755 --- a/scripts/stress-test/setup/login_admin.py +++ b/scripts/stress-test/setup/login_admin.py @@ -5,13 +5,19 @@ from pathlib import Path sys.path.append(str(Path(__file__).parent.parent)) +import base64 import json import httpx from common import Logger, config_helper -def login_admin() -> None: +def encode_sensitive_field(value: str) -> str: + """Encode fields the same way the web client does before login.""" + return base64.b64encode(value.encode("utf-8")).decode() + + +def login_admin() -> bool: """Login with admin account and save access token.""" log = Logger("Login") @@ -23,7 +29,7 @@ def login_admin() -> None: if not admin_config: log.error("Admin config not found") log.info("Please run setup_admin.py first to create the admin account") - return + return False log.info(f"Logging in with email: {admin_config['email']}") @@ -34,7 +40,7 @@ def login_admin() -> None: # Prepare login payload login_payload = { "email": admin_config["email"], - "password": admin_config["password"], + "password": encode_sensitive_field(admin_config["password"]), "remember_me": True, } @@ -50,29 +56,28 @@ def login_admin() -> None: if response.status_code == 200: log.success("Login successful!") - # Extract token from response response_data = response.json() # Check if login was successful if response_data.get("result") != "success": log.error(f"Login failed: {response_data}") - return + return False - # Extract tokens from data field - token_data = response_data.get("data", {}) - access_token = token_data.get("access_token", "") - refresh_token = token_data.get("refresh_token", "") + access_token = response.cookies.get("access_token", "") + refresh_token = response.cookies.get("refresh_token", "") + csrf_token = response.cookies.get("csrf_token", "") if not access_token: log.error("No access token found in response") log.debug(f"Full response: {json.dumps(response_data, indent=2)}") - return + return False # Save token to config file token_config = { "email": admin_config["email"], "access_token": access_token, "refresh_token": refresh_token, + "csrf_token": csrf_token, } # Save token config @@ -82,20 +87,26 @@ def login_admin() -> None: # Show truncated token for verification token_display = f"{access_token[:20]}..." if len(access_token) > 20 else "Token saved" log.key_value("Access token", token_display) + return True elif response.status_code == 401: log.error("Login failed: Invalid credentials") log.debug(f"Response: {response.text}") + return False else: log.error(f"Login failed with status code: {response.status_code}") log.debug(f"Response: {response.text}") + return False except httpx.ConnectError: log.error("Could not connect to Dify API at http://localhost:5001") log.info("Make sure the API server is running with: ./dev/start-api") + return False except Exception as e: log.error(f"An error occurred: {e}") + return False if __name__ == "__main__": - login_admin() + if not login_admin(): + sys.exit(1) diff --git a/scripts/stress-test/setup/publish_workflow.py b/scripts/stress-test/setup/publish_workflow.py index b772eccebdd..3e565ff7624 100755 --- a/scripts/stress-test/setup/publish_workflow.py +++ b/scripts/stress-test/setup/publish_workflow.py @@ -11,7 +11,7 @@ import httpx from common import Logger, config_helper -def publish_workflow() -> None: +def publish_workflow() -> bool: """Publish the imported workflow app.""" log = Logger("PublishWorkflow") @@ -21,13 +21,13 @@ def publish_workflow() -> None: access_token = config_helper.get_token() if not access_token: log.error("No access token found in config") - return + return False # Read app_id from config app_id = config_helper.get_app_id() if not app_id: log.error("No app_id found in config") - return + return False log.step(f"Publishing workflow for app: {app_id}") @@ -51,14 +51,14 @@ def publish_workflow() -> None: "Sec-Fetch-Mode": "cors", "Sec-Fetch-Site": "same-site", "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/139.0.0.0 Safari/537.36", - "authorization": f"Bearer {access_token}", + **config_helper.console_auth_headers(), "content-type": "application/json", "sec-ch-ua": '"Not;A=Brand";v="99", "Google Chrome";v="139", "Chromium";v="139"', "sec-ch-ua-mobile": "?0", "sec-ch-ua-platform": '"macOS"', } - cookies = {"locale": "en-US"} + cookies = config_helper.console_auth_cookies() try: # Make the publish request @@ -83,23 +83,30 @@ def publish_workflow() -> None: except json.JSONDecodeError: # Response might be empty or non-JSON pass + return True elif response.status_code == 401: log.error("Workflow publish failed: Unauthorized") log.info("Token may have expired. Please run login_admin.py again") + return False elif response.status_code == 404: log.error("Workflow publish failed: App not found") log.info("Make sure the app was imported successfully") + return False else: log.error(f"Workflow publish failed with status code: {response.status_code}") log.debug(f"Response: {response.text}") + return False except httpx.ConnectError: log.error("Could not connect to Dify API at http://localhost:5001") log.info("Make sure the API server is running with: ./dev/start-api") + return False except Exception as e: log.error(f"An error occurred: {e}") + return False if __name__ == "__main__": - publish_workflow() + if not publish_workflow(): + sys.exit(1) diff --git a/scripts/stress-test/setup/setup_admin.py b/scripts/stress-test/setup/setup_admin.py index a5e9161210c..6e5dbd373f9 100755 --- a/scripts/stress-test/setup/setup_admin.py +++ b/scripts/stress-test/setup/setup_admin.py @@ -5,22 +5,27 @@ from pathlib import Path sys.path.append(str(Path(__file__).parent.parent)) +import os + import httpx from common import Logger, config_helper -def setup_admin_account() -> None: +def build_admin_config() -> dict[str, str]: + return { + "email": os.getenv("STRESS_TEST_ADMIN_EMAIL", "test@dify.ai"), + "username": os.getenv("STRESS_TEST_ADMIN_USERNAME", "dify"), + "password": os.getenv("STRESS_TEST_ADMIN_PASSWORD", "password123"), + } + + +def setup_admin_account() -> bool: """Setup Dify API with an admin account.""" log = Logger("SetupAdmin") log.header("Setting up Admin Account") - # Admin account credentials - admin_config = { - "email": "test@dify.ai", - "username": "dify", - "password": "password123", - } + admin_config = build_admin_config() # Save credentials to config file if config_helper.write_config("admin_config", admin_config): @@ -52,20 +57,26 @@ def setup_admin_account() -> None: log.success("Admin account created successfully!") log.key_value("Email", admin_config["email"]) log.key_value("Username", admin_config["username"]) + return True - elif response.status_code == 400: - log.warning("Setup may have already been completed or invalid data provided") + elif response.status_code in {400, 403}: + log.warning("Setup may have already been completed") log.debug(f"Response: {response.text}") + return True else: log.error(f"Setup failed with status code: {response.status_code}") log.debug(f"Response: {response.text}") + return False except httpx.ConnectError: log.error("Could not connect to Dify API at http://localhost:5001") log.info("Make sure the API server is running with: ./dev/start-api") + return False except Exception as e: log.error(f"An error occurred: {e}") + return False if __name__ == "__main__": - setup_admin_account() + if not setup_admin_account(): + sys.exit(1) diff --git a/scripts/stress-test/setup_all.py b/scripts/stress-test/setup_all.py index ece420f9257..47a789c4317 100755 --- a/scripts/stress-test/setup_all.py +++ b/scripts/stress-test/setup_all.py @@ -1,12 +1,56 @@ #!/usr/bin/env python3 +import os import socket import subprocess import sys import time from pathlib import Path -from common import Logger, ProgressLogger +import httpx +from common import Logger, ProgressLogger, config_helper + + +def build_admin_config() -> dict[str, str]: + return { + "email": os.getenv("STRESS_TEST_ADMIN_EMAIL", "test@dify.ai"), + "username": os.getenv("STRESS_TEST_ADMIN_USERNAME", "dify"), + "password": os.getenv("STRESS_TEST_ADMIN_PASSWORD", "password123"), + } + + +def get_setup_step(base_url: str = "http://localhost:5001") -> str | None: + try: + response = httpx.get(f"{base_url}/console/api/setup", timeout=5) + if response.status_code == 200: + return response.json().get("step") + except (httpx.HTTPError, ValueError): + return None + return None + + +def confirm(prompt: str) -> bool: + answer = input(f"\n{prompt} [Y/n]: ").strip().lower() + return answer in ("", "y", "yes") + + +def confirm_admin_credentials(log: Logger) -> bool: + admin_config = build_admin_config() + setup_step = get_setup_step() + config_helper.write_config("admin_config", admin_config) + + if setup_step == "finished": + log.warning("Dify is already initialized; setup will use the existing admin account to log in.") + log.key_value("Admin email", admin_config["email"]) + log.info("Set STRESS_TEST_ADMIN_EMAIL and STRESS_TEST_ADMIN_PASSWORD if this is not the right account.") + return confirm("Continue with this admin login?") + + log.info("Dify is not initialized; setup will create the first admin account with:") + log.key_value("Admin email", admin_config["email"]) + log.key_value("Admin username", admin_config["username"]) + log.key_value("Admin password", admin_config["password"]) + log.info("Set STRESS_TEST_ADMIN_EMAIL, STRESS_TEST_ADMIN_USERNAME, or STRESS_TEST_ADMIN_PASSWORD to override.") + return confirm("Create/use this admin account?") def run_script(script_name: str, description: str) -> bool: @@ -89,19 +133,20 @@ def main() -> None: if not dify_running or not mock_running: print("\n⚠️ Both services must be running before proceeding.") - retry = input("\nWould you like to check again? (yes/no): ") - if retry.lower() in ["yes", "y"]: + if confirm("Would you like to check again?"): return main() # Recursively call main to check again else: print("❌ Setup cancelled. Please start the required services and try again.") sys.exit(1) log.success("All required services are running!") - input("\nPress Enter to continue with setup...") + if not confirm_admin_credentials(log): + print("❌ Setup cancelled. Please set the admin environment variables and try again.") + sys.exit(1) # Define setup steps + setup_step = get_setup_step() setup_steps = [ - ("setup_admin.py", "Creating admin account"), ("login_admin.py", "Logging in and getting access token"), ("install_openai_plugin.py", "Installing OpenAI plugin"), ("configure_openai_plugin.py", "Configuring OpenAI plugin with mock server"), @@ -109,6 +154,8 @@ def main() -> None: ("create_api_key.py", "Creating API key for the app"), ("publish_workflow.py", "Publishing the workflow"), ] + if setup_step != "finished": + setup_steps.insert(0, ("setup_admin.py", "Creating admin account")) # Create progress logger progress = ProgressLogger(len(setup_steps), log) @@ -131,6 +178,18 @@ def main() -> None: log.error(f"Setup failed at: {failed_step}") log.separator() log.info("Troubleshooting:") + if failed_step == "login_admin.py": + if get_setup_step() == "finished": + log.list_item( + "Dify is already initialized; set STRESS_TEST_ADMIN_EMAIL and " + "STRESS_TEST_ADMIN_PASSWORD to an existing admin account." + ) + else: + admin_config = build_admin_config() + log.list_item( + "Dify is not initialized; setup creates the first admin with " + f"{admin_config['email']} / {admin_config['username']} unless overridden by environment variables." + ) log.list_item("Check if the Dify API server is running (./dev/start-api)") log.list_item("Check if the mock OpenAI server is running (port 5004)") log.list_item("Review the error messages above") @@ -151,9 +210,7 @@ def main() -> None: # Optionally run a test log.separator() - test_input = input("Would you like to run a test workflow now? (yes/no): ") - - if test_input.lower() in ["yes", "y"]: + if confirm("Would you like to run a test workflow now?"): log.step("Running test workflow...") run_script("run_workflow.py", "Testing workflow with default question") diff --git a/scripts/stress-test/test_setup_scripts.py b/scripts/stress-test/test_setup_scripts.py new file mode 100644 index 00000000000..01ddd2b8fb9 --- /dev/null +++ b/scripts/stress-test/test_setup_scripts.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 + +import base64 +import importlib.util +import sys +from pathlib import Path + +from common.config_helper import ConfigHelper + + +def _load_setup_module(name: str): + module_path = Path(__file__).parent / "setup" / f"{name}.py" + spec = importlib.util.spec_from_file_location(f"stress_test_{name}", module_path) + assert spec and spec.loader + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _load_stress_test_module(name: str): + module_path = Path(__file__).parent / f"{name}.py" + spec = importlib.util.spec_from_file_location(f"stress_test_{name}", module_path) + assert spec and spec.loader + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def test_config_helper_getters_read_state_sections(tmp_path): + helper = ConfigHelper(base_dir=tmp_path) + helper.write_state( + { + "auth": {"access_token": "console-token", "csrf_token": "csrf-token"}, + "app": {"app_id": "app-id"}, + "api_key": {"token": "app-token"}, + } + ) + + assert helper.get_token() == "console-token" + assert helper.get_csrf_token() == "csrf-token" + assert helper.get_app_id() == "app-id" + assert helper.get_api_key() == "app-token" + assert helper.console_auth_headers() == { + "authorization": "Bearer console-token", + "X-CSRF-Token": "csrf-token", + } + assert helper.console_auth_cookies() == { + "locale": "en-US", + "access_token": "console-token", + "csrf_token": "csrf-token", + } + + +def test_login_admin_encodes_password_like_web_client(): + login_admin = _load_setup_module("login_admin") + + encoded = login_admin.encode_sensitive_field("password123") + + assert encoded == base64.b64encode(b"password123").decode() + + +def test_setup_admin_reads_credentials_from_environment(monkeypatch): + setup_admin = _load_setup_module("setup_admin") + monkeypatch.setenv("STRESS_TEST_ADMIN_EMAIL", "real-admin@example.com") + monkeypatch.setenv("STRESS_TEST_ADMIN_USERNAME", "real-admin") + monkeypatch.setenv("STRESS_TEST_ADMIN_PASSWORD", "secret") + + assert setup_admin.build_admin_config() == { + "email": "real-admin@example.com", + "username": "real-admin", + "password": "secret", + } + + +def test_setup_all_reads_credentials_from_environment(monkeypatch): + setup_all = _load_stress_test_module("setup_all") + monkeypatch.setenv("STRESS_TEST_ADMIN_EMAIL", "real-admin@example.com") + monkeypatch.setenv("STRESS_TEST_ADMIN_USERNAME", "real-admin") + monkeypatch.setenv("STRESS_TEST_ADMIN_PASSWORD", "secret") + + assert setup_all.build_admin_config() == { + "email": "real-admin@example.com", + "username": "real-admin", + "password": "secret", + } + + +def test_setup_all_confirm_defaults_to_yes(monkeypatch): + setup_all = _load_stress_test_module("setup_all") + monkeypatch.setattr("builtins.input", lambda prompt: "") + + assert setup_all.confirm("Continue?") is True + + +def test_create_api_key_fails_without_app_id(tmp_path): + create_api_key = _load_setup_module("create_api_key") + create_api_key.config_helper.base_dir = tmp_path + create_api_key.config_helper.write_state({"auth": {"access_token": "token"}}) + + assert create_api_key.create_api_key() is False + + +def test_plugin_install_response_without_task_is_non_blocking(): + install_openai_plugin = _load_setup_module("install_openai_plugin") + + assert install_openai_plugin.is_non_blocking_install_response({"ok": True}) is True + assert install_openai_plugin.is_non_blocking_install_response({}) is True + assert install_openai_plugin.is_non_blocking_install_response({"code": "plugin_error"}) is False + + +def test_import_response_with_warnings_and_app_id_is_success(): + import_workflow_app = _load_setup_module("import_workflow_app") + + assert import_workflow_app.is_successful_import_response( + {"status": "completed-with-warnings", "app_id": "app-id"} + ) + assert not import_workflow_app.is_successful_import_response({"status": "failed", "app_id": "app-id"}) + assert not import_workflow_app.is_successful_import_response({"status": "completed"})