package main import ( "context" "encoding/json" "flag" "fmt" "net/http" "os" "os/signal" "sync" "syscall" "time" "plp-test/internal/config" "plp-test/internal/model" "plp-test/internal/testcase" "plp-test/internal/utils" "github.com/sirupsen/logrus" ) var ( configFile string logLevel string ) func init() { flag.StringVar(&configFile, "config", "config.yaml", "配置文件路径") flag.StringVar(&logLevel, "log-level", "info", "日志级别 (debug, info, warn, error)") } // TestRunner 测试运行器 type TestRunner struct { config *config.Config logger *logrus.Logger factory *testcase.TestCaseFactory tests map[string]testcase.TestCase testsMu sync.RWMutex testResult map[string]*model.TestResult resultMu sync.RWMutex } // NewTestRunner 创建测试运行器 func NewTestRunner(cfg *config.Config, logger *logrus.Logger) *TestRunner { return &TestRunner{ config: cfg, logger: logger, factory: testcase.NewTestCaseFactory(cfg, logger), tests: make(map[string]testcase.TestCase), testResult: make(map[string]*model.TestResult), } } // RunTest 运行指定的测试 func (r *TestRunner) RunTest(testType string) (*model.TestResult, error) { r.logger.Infof("准备运行测试: %s", testType) // 创建测试实例 test, err := r.factory.CreateTestCase(testType) if err != nil { return nil, fmt.Errorf("创建测试用例失败: %v", err) } if test == nil { return nil, fmt.Errorf("未找到测试用例: %s", testType) } // 存储测试实例 testID := test.Status().TestID r.testsMu.Lock() r.tests[testID] = test r.testsMu.Unlock() // 创建上下文以便可以取消测试 ctx, cancel := context.WithCancel(context.Background()) defer cancel() // 设置测试环境 r.logger.Info("设置测试环境") if err := test.Setup(ctx); err != nil { r.logger.Errorf("设置测试环境失败: %v", err) return nil, err } // 运行测试 r.logger.Info("运行测试") result, err := test.Run(ctx) if err != nil { r.logger.Errorf("测试运行失败: %v", err) // 尝试清理 cleanupErr := test.Cleanup(ctx) if cleanupErr != nil { r.logger.Errorf("测试清理失败: %v", cleanupErr) } return nil, err } // 清理测试环境 r.logger.Info("清理测试环境") if err := test.Cleanup(ctx); err != nil { r.logger.Errorf("测试清理失败: %v", err) return nil, err } // 存储测试结果 r.resultMu.Lock() r.testResult[testID] = result r.resultMu.Unlock() // 移除测试实例 r.testsMu.Lock() delete(r.tests, testID) r.testsMu.Unlock() r.logger.Infof("测试 %s 完成", testType) return result, nil } // GetTestStatus 获取测试状态 func (r *TestRunner) GetTestStatus(testID string) *model.TestStatus { r.testsMu.RLock() defer r.testsMu.RUnlock() if test, ok := r.tests[testID]; ok { return test.Status() } return nil } // GetAllTestStatus 获取所有测试状态 func (r *TestRunner) GetAllTestStatus() []*model.TestStatus { r.testsMu.RLock() defer r.testsMu.RUnlock() statuses := make([]*model.TestStatus, 0, len(r.tests)) for _, test := range r.tests { statuses = append(statuses, test.Status()) } return statuses } // StartServer 启动HTTP服务器 func StartServer(cfg *config.Config, runner *TestRunner, logger *logrus.Logger) *http.Server { mux := http.NewServeMux() // 健康检查接口 mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) { health := &model.HealthStatus{ Status: "ok", Timestamp: time.Now(), Message: "服务正常运行", } w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(health) }) // 运行测试接口 mux.HandleFunc("/run", func(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodPost { http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) return } var req model.TestRequest if err := json.NewDecoder(r.Body).Decode(&req); err != nil { http.Error(w, "Invalid request body", http.StatusBadRequest) return } logger.Infof("收到测试请求: %+v", req) // 异步运行测试 go func() { result, err := runner.RunTest(req.TestType) if err != nil { logger.Errorf("测试运行失败: %v", err) } else { logger.Infof("测试完成: %+v", result) } }() resp := model.TestResponse{ RequestID: req.TestType + "-" + time.Now().Format("20060102-150405"), Status: "accepted", Message: "测试已接受并开始执行", ServerTime: time.Now(), } w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusAccepted) json.NewEncoder(w).Encode(resp) }) // 获取测试状态接口 mux.HandleFunc("/status", func(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodGet { http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) return } testID := r.URL.Query().Get("test_id") var status interface{} if testID == "" { // 获取所有测试状态 status = runner.GetAllTestStatus() } else { // 获取指定测试状态 status = runner.GetTestStatus(testID) if status == nil { http.Error(w, "Test not found", http.StatusNotFound) return } } w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(status) }) // 启动服务器 addr := fmt.Sprintf("%s:%d", cfg.Server.ListenAddr, cfg.Server.Port) server := &http.Server{ Addr: addr, Handler: mux, } go func() { logger.Infof("服务器启动在 %s", addr) if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed { logger.Fatalf("服务器启动失败: %v", err) } }() return server } func main() { flag.Parse() // 初始化日志级别 var level logrus.Level switch logLevel { case "debug": level = logrus.DebugLevel case "info": level = logrus.InfoLevel case "warn": level = logrus.WarnLevel case "error": level = logrus.ErrorLevel default: level = logrus.InfoLevel } // 初始化日志 logger := logrus.New() logger.SetLevel(level) logger.SetFormatter(&logrus.TextFormatter{ FullTimestamp: true, TimestampFormat: "2006-01-02 15:04:05", }) // 加载配置 logger.Infof("加载配置文件: %s", configFile) cfg, err := config.Load(configFile) if err != nil { logger.Fatalf("加载配置失败: %v", err) } // 初始化日志文件 if cfg.Server.LogFile != "" { utils.InitLogger(cfg.Server.LogFile, level) logger = utils.Logger } // 创建测试运行器 runner := NewTestRunner(cfg, logger) // 启动服务器 server := StartServer(cfg, runner, logger) // 等待终止信号 stop := make(chan os.Signal, 1) signal.Notify(stop, os.Interrupt, syscall.SIGTERM) <-stop logger.Info("正在关闭服务器...") ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() if err := server.Shutdown(ctx); err != nil { logger.Fatalf("服务器强制关闭: %v", err) } logger.Info("服务器已优雅关闭") }