Files
dify/scripts/stress-test/setup/setup_admin.py

83 lines
2.5 KiB
Python
Executable File

#!/usr/bin/env python3
import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent))
import os
import httpx
from common import Logger, 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 setup_admin_account() -> bool:
"""Setup Dify API with an admin account."""
log = Logger("SetupAdmin")
log.header("Setting up Admin Account")
admin_config = build_admin_config()
# Save credentials to config file
if config_helper.write_config("admin_config", admin_config):
log.info(f"Admin credentials saved to: {config_helper.get_config_path('benchmark_state')}")
# API setup endpoint
base_url = "http://localhost:5001"
setup_endpoint = f"{base_url}/console/api/setup"
# Prepare setup payload
setup_payload = {
"email": admin_config["email"],
"name": admin_config["username"],
"password": admin_config["password"],
}
log.step("Configuring Dify with admin account...")
try:
# Make the setup request
with httpx.Client() as client:
response = client.post(
setup_endpoint,
json=setup_payload,
headers={"Content-Type": "application/json"},
)
if response.status_code == 201:
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 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__":
if not setup_admin_account():
sys.exit(1)