Coverage for main.py: 63.89%

1from contextlib import asynccontextmanager 

2 

3import uvicorn 

4from fastapi import FastAPI 

5from fastapi.middleware.cors import CORSMiddleware 

6 

7from api.v1.endpoints import (api_test, assets, auth, backtest, broker, 

8 risk_management) 

9from api.v1.endpoints import settings as settings_endpoints 

10from api.v1.endpoints import (stock, stock_import, strategies, trade_test, 

11 trading, users, websocket) 

12from core.services.auth_service import AuthService 

13from infrastructure.config.settings import settings 

14from infrastructure.database.redis_client import get_redis 

15 

16 

17@asynccontextmanager 

18async def lifespan(app: FastAPI): 

19 """应用生命周期管理""" 

20 # 启动时创建默认管理员用户 

21 auth_service = AuthService() 

22 auth_service.create_default_admin() 

23 yield 

24 # 关闭时的清理工作（如果需要） 

25 

26 

27app = FastAPI( 

28 title=settings.app_name, 

29 description="用于量化交易股票的后端API服务", 

30 version=settings.app_version, 

31 lifespan=lifespan, 

32) 

33 

34# 配置CORS 

35app.add_middleware( 

36 CORSMiddleware, 

37 allow_origins=["http://localhost:3000", "http://localhost:5173"], # 前端开发服务器 

38 allow_credentials=True, 

39 allow_methods=["*"], 

40 allow_headers=["*"], 

41) 

42 

43 

44@app.get("/") 

45async def root(): 

46 return {"message": "量化交易系统 API 服务正在运行"} 

47 

48 

49@app.get("/health") 

50async def health_check(): 

51 """健康检查端点""" 

52 redis_client = get_redis() 

53 redis_status = "connected" if redis_client.is_connected() else "disconnected" 

54 

55 return { 

56 "status": "healthy", 

57 "redis": redis_status, 

58 "services": {"api": "running", "redis": redis_status}, 

59 } 

60 

61 

62@app.get("/redis/status") 

63async def redis_status(): 

64 """Redis状态检查""" 

65 redis_client = get_redis() 

66 

67 if not redis_client.is_connected(): 

68 return {"status": "disconnected", "error": "Redis连接失败"} 

69 

70 try: 

71 info = redis_client.info() 

72 return { 

73 "status": "connected", 

74 "info": { 

75 "version": info.get("redis_version", "unknown"), 

76 "uptime": info.get("uptime_in_seconds", 0), 

77 "connected_clients": info.get("connected_clients", 0), 

78 "used_memory": info.get("used_memory_human", "0B"), 

79 }, 

80 } 

81 except Exception as e: 

82 return {"status": "error", "error": str(e)} 

83 

84 

85@app.get("/redis/test") 

86async def redis_test(): 

87 """Redis测试端点""" 

88 redis_client = get_redis() 

89 

90 if not redis_client.is_connected(): 

91 return {"status": "error", "message": "Redis连接失败"} 

92 

93 try: 

94 # 测试设置和获取 

95 test_key = "test:connection" 

96 test_value = {"message": "Redis连接测试", "timestamp": "2024-01-01T00:00:00Z"} 

97 

98 # 设置测试数据 

99 redis_client.set(test_key, test_value, expire=60) 

100 

101 # 获取测试数据 

102 retrieved_value = redis_client.get(test_key) 

103 

104 # 清理测试数据 

105 redis_client.delete(test_key) 

106 

107 return { 

108 "status": "success", 

109 "message": "Redis连接和操作正常", 

110 "test_data": retrieved_value, 

111 } 

112 except Exception as e: 

113 return {"status": "error", "message": f"Redis测试失败: {str(e)}"} 

114 

115 

116# 添加API路由 

117app.include_router(stock.router, prefix="/api/v1") 

118app.include_router(auth.router, prefix="/api/v1") 

119app.include_router(users.router, prefix="/api/v1") 

120app.include_router(broker.router, prefix="/api/v1") 

121app.include_router(settings_endpoints.router, prefix="/api/v1") 

122app.include_router(stock_import.router, prefix="/api/v1/stock-import") 

123app.include_router(assets.router, prefix="/api/v1") 

124# app.include_router(unified_assets.router, prefix="/api/v1/unified-assets") # 已删除 

125app.include_router(api_test.router, prefix="/api/v1") 

126app.include_router(trade_test.router, prefix="/api/v1") 

127app.include_router(trading.router, prefix="/api/v1/trading") 

128app.include_router(strategies.router, prefix="/api/v1/strategies") 

129app.include_router(backtest.router, prefix="/api/v1/backtest") 

130app.include_router(risk_management.router, prefix="/api/v1/risk") 

131 

132# 添加WebSocket路由 

133from fastapi import WebSocket 

134 

135from core.services.websocket_service import websocket_endpoint 

136 

137 

138@app.websocket("/ws/stock-import/{task_id}") 

139async def websocket_endpoint_handler(websocket: WebSocket, task_id: str): 

140 await websocket_endpoint(websocket, task_id) 

141 

142 

143@app.websocket("/ws/trading/{session_id}") 

144async def trading_websocket_endpoint(websocket: WebSocket, session_id: str): 

145 await websocket_endpoint(websocket, session_id) 

146 

147 

148# 启动事件已移至lifespan函数中 

149 

150if __name__ == "__main__": 

151 uvicorn.run(app, host="0.0.0.0", port=3000)