stress-test-kit 0.1.0__py3-none-any.whl

stress_test_kit/__init__.py

@@ -0,0 +1,5 @@
+"""
+Stress Test Kit - 本地压测工具
+"""
+
+__version__ = "0.1.0"

stress_test_kit/__main__.py

@@ -0,0 +1,5 @@
+"""CLI 入口模块"""
+from stress_test_kit.cli import main
+
+if __name__ == "__main__":
+    main()

stress_test_kit/cli.py

@@ -0,0 +1,115 @@
+"""CLI 入口（类似 locust 命令）"""
+import os
+import click
+import uvicorn
+import webbrowser
+import threading
+
+
+@click.command()
+@click.option('--port', default=8765, help='Web服务端口（前后端同一端口，不推荐与 --backend-port/--frontend-port 同时使用）')
+@click.option('--backend-port', default=None, type=int, help='后端 API 服务端口（优先级高于 --port）')
+@click.option('--frontend-port', default=None, type=int, help='前端静态服务端口（优先级高于 --port，需与 --backend-port 同时指定）')
+@click.option('--data-dir', default=None, help='数据目录（存放SQLite和文件）')
+@click.option('--no-browser', is_flag=True, help='不自动打开浏览器')
+def main(port: int, backend_port: int, frontend_port: int, data_dir: str, no_browser: bool):
+    """启动 Stress Test Kit 本地压测服务"""
+    import sys
+    from stress_test_kit.process_manager import _is_port_free
+
+    # 端口配置逻辑
+    if backend_port is not None and frontend_port is not None:
+        # 前后端分离模式
+        actual_backend_port = backend_port
+        actual_frontend_port = frontend_port
+        frontend_separate = True
+    elif backend_port is not None or frontend_port is not None:
+        # 只指定了一个，报错
+        click.echo("❌ --backend-port 和 --frontend-port 必须同时指定", err=True)
+        sys.exit(1)
+    else:
+        # 使用 --port（前后端同一端口）
+        actual_backend_port = port
+        actual_frontend_port = None
+        frontend_separate = False
+
+    # 端口占用检查
+    if not _is_port_free(actual_backend_port):
+        click.echo(f"❌ 后端端口 {actual_backend_port} 已被占用，请先释放该端口或指定其他端口", err=True)
+        sys.exit(1)
+    
+    if frontend_separate and not _is_port_free(actual_frontend_port):
+        click.echo(f"❌ 前端端口 {actual_frontend_port} 已被占用，请先释放该端口或指定其他端口", err=True)
+        sys.exit(1)
+
+    # 数据目录默认值
+    if data_dir is None:
+        data_dir = os.path.join(os.path.expanduser("~"), ".stress-test-kit")
+    else:
+        data_dir = os.path.expanduser(data_dir)
+    
+    # 初始化数据库
+    from stress_test_kit.models import init_db
+    from stress_test_kit.server import set_data_dir, set_mount_frontend
+    from stress_test_kit.sampler import recover_active_samplers, start_periodic_patrol
+    
+    init_db(data_dir)
+    set_data_dir(data_dir)
+    
+    # 前后端分离模式下，禁用 FastAPI 挂载静态文件
+    if frontend_separate:
+        set_mount_frontend(False)
+    
+    # 恢复活跃的采样线程
+    recover_active_samplers(data_dir)
+    
+    # 启动定期巡检（每 60 秒检查所有 running 状态的任务）
+    start_periodic_patrol(data_dir)
+    
+    click.echo("=" * 60)
+    click.echo("🚀 Stress Test Kit 启动中...")
+    click.echo("=" * 60)
+    click.echo(f"📂 数据目录: {data_dir}")
+    
+    if frontend_separate:
+        click.echo(f"🔧 后端 API: http://localhost:{actual_backend_port}")
+        click.echo(f"🌐 前端页面: http://localhost:{actual_frontend_port}")
+        access_url = f"http://localhost:{actual_frontend_port}"
+        
+        # 启动前端静态文件服务（独立线程）
+        def start_frontend_server():
+            import http.server
+            import socketserver
+            import functools
+            frontend_dir = os.path.join(os.path.dirname(__file__), "frontend")
+            if os.path.exists(frontend_dir):
+                os.chdir(frontend_dir)
+                handler = functools.partial(http.server.SimpleHTTPRequestHandler, directory=frontend_dir)
+                with socketserver.TCPServer(("0.0.0.0", actual_frontend_port), handler) as httpd:
+                    click.echo(f"✅ 前端服务已启动在端口 {actual_frontend_port}")
+                    httpd.serve_forever()
+            else:
+                click.echo(f"⚠️  前端目录不存在: {frontend_dir}，跳过前端服务启动")
+        
+        threading.Thread(target=start_frontend_server, daemon=True).start()
+    else:
+        click.echo(f"🌐 访问地址: http://localhost:{actual_backend_port}")
+        access_url = f"http://localhost:{actual_backend_port}"
+    
+    click.echo("=" * 60)
+    
+    # 自动打开浏览器
+    if not no_browser:
+        threading.Timer(1.5, lambda: webbrowser.open(access_url)).start()
+    
+    # 启动 FastAPI 服务
+    uvicorn.run(
+        "stress_test_kit.server:app",
+        host="0.0.0.0",
+        port=actual_backend_port,
+        log_level="info"
+    )
+
+
+if __name__ == "__main__":
+    main()