plp-test/cmd/server/main.go

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("服务器已优雅关闭")
}