4. Go testing 高级技巧：TestMain、Setup/Teardown 与 Test Helpers

在 Go 项目的压测与集成测试中，仅靠基础的 TestXxx 函数往往无法满足复杂的资源管理需求。本章将深入 testing 包的三个高级模式：TestMain、Setup/Teardown 以及 Test Helpers，帮助你在 Kafka 压测场景中构建健壮、可复用的测试基础设施。

4.1 TestMain：测试入口的全局控制

TestMain 是 testing 包提供的特殊函数，它允许你在测试运行前后执行自定义逻辑。当你在包中定义 func TestMain(m *testing.M) 时，Go 测试框架会调用它而非直接运行测试用例。这是管理全局资源（如 Kafka 集群连接、数据库初始化、环境变量设置）的理想位置。

func TestMain(m *testing.M) {
    // 1. 全局初始化：连接 Kafka 集群
    brokers := []string{"localhost:9092"}
    adminClient, err := kafka.NewAdminClient(&kafka.ConfigMap{
        "bootstrap.servers": brokers[0],
    })
    if err != nil {
        log.Fatalf("Failed to create admin client: %v", err)
    }
    defer adminClient.Close()

    // 2. 创建测试 Topic
    topic := "benchmark-topic"
    _, err = adminClient.CreateTopics(context.Background(),
        []kafka.TopicSpecification{{
            Topic:         topic,
            NumPartitions: 6,
        }})
    if err != nil {
        log.Fatalf("Failed to create topic: %v", err)
    }

    // 3. 运行所有测试
    exitCode := m.Run()

    // 4. 全局清理：删除测试 Topic
    adminClient.DeleteTopics(context.Background(), []string{topic})

    // 5. 退出
    os.Exit(exitCode)
}

4.2 Setup 与 Teardown 模式：细粒度资源管理

虽然 TestMain 提供了全局控制，但许多场景需要更细粒度的资源管理：每个测试用例或每个测试组独立初始化与清理。Go 标准库没有内置的 Setup/Teardown 钩子，但我们可以通过 t.Cleanup 和嵌套函数优雅地实现这一模式。

4.2.1 基于 t.Cleanup 的 Setup/Teardown

func TestKafkaProducer(t *testing.T) {
    // Setup：创建生产者
    producer, err := kafka.NewProducer(&kafka.ConfigMap{
        "bootstrap.servers": "localhost:9092",
    })
    if err != nil {
        t.Fatalf("Failed to create producer: %v", err)
    }

    // Teardown：注册清理函数
    t.Cleanup(func() {
        producer.Flush(15 * 1000)
        producer.Close()
        t.Log("Producer closed successfully")
    })

    // 测试逻辑
    topic := "benchmark-topic"
    err = producer.Produce(&kafka.Message{
        TopicPartition: kafka.TopicPartition{Topic: &topic, Partition: kafka.PartitionAny},
        Value:          []byte("test message"),
    }, nil)

    if err != nil {
        t.Errorf("Produce failed: %v", err)
    }
}

4.2.2 测试组级别的 Setup/Teardown

当多个测试用例共享相同的初始化逻辑时，可以使用 t.Run 创建子测试组，并在组级别统一管理资源。

func TestKafkaConsumerGroup(t *testing.T) {
    // 组级别 Setup
    brokers := []string{"localhost:9092"}
    groupID := "test-group"

    // 为整个组创建消费者
    consumer, err := kafka.NewConsumer(&kafka.ConfigMap{
        "bootstrap.servers": brokers[0],
        "group.id":         groupID,
        "auto.offset.reset": "earliest",
    })
    if err != nil {
        t.Fatalf("Failed to create consumer: %v", err)
    }

    // 组级别 Teardown
    t.Cleanup(func() {
        consumer.Close()
        t.Log("Consumer group closed")
    })

    // 子测试 1：订阅 Topic
    t.Run("Subscribe", func(t *testing.T) {
        err := consumer.SubscribeTopics([]string{"benchmark-topic"}, nil)
        if err != nil {
            t.Errorf("Subscribe failed: %v", err)
        }
    })

    // 子测试 2：消费消息
    t.Run("Consume", func(t *testing.T) {
        msg, err := consumer.ReadMessage(1000)
        if err != nil {
            t.Errorf("ReadMessage failed: %v", err)
        }
        t.Logf("Received message: %s", string(msg.Value))
    })
}

4.3 Test Helpers：消除样板代码

在 Kafka 压测中，创建生产者、消费者、Topic 等操作频繁重复。Test Helpers 是返回 testing.TB 参数的函数，用于封装这些重复逻辑，同时自动处理错误和清理。按照 Go 惯例，helper 函数应调用 t.Helper() 标记自身，使错误报告指向调用者而非 helper 内部。

// createTestProducer 创建一个测试用 Kafka 生产者，并注册清理
func createTestProducer(t testing.TB) *kafka.Producer {
    t.Helper() // 标记为 helper 函数

    p, err := kafka.NewProducer(&kafka.ConfigMap{
        "bootstrap.servers": "localhost:9092",
    })
    if err != nil {
        t.Fatalf("createTestProducer: %v", err)
    }

    t.Cleanup(func() {
        p.Flush(15 * 1000)
        p.Close()
    })

    return p
}

// createTestTopic 创建测试 Topic，返回 Topic 名称
func createTestTopic(t testing.TB, name string, partitions int) string {
    t.Helper()

    admin, err := kafka.NewAdminClient(&kafka.ConfigMap{
        "bootstrap.servers": "localhost:9092",
    })
    if err != nil {
        t.Fatalf("createTestTopic admin client: %v", err)
    }
    defer admin.Close()

    _, err = admin.CreateTopics(context.Background(),
        []kafka.TopicSpecification{{
            Topic:         name,
            NumPartitions: partitions,
        }})
    if err != nil {
        t.Fatalf("createTestTopic: %v", err)
    }

    t.Cleanup(func() {
        admin.DeleteTopics(context.Background(), []string{name})
    })

    return name
}

4.4 综合架构：压测闭环中的高级模式

将上述三种模式组合，可以构建一个完整的 Kafka 压测闭环。下面的 SVG 图展示了 TestMain、Setup/Teardown 和 Test Helpers 在测试生命周期中的协作关系。

4.5 最佳实践与避坑指南

• TestMain 中避免使用 t 对象：TestMain 接收 *testing.M 而非 *testing.T，无法调用 t.Helper() 或 t.Cleanup()。全局清理应使用 defer 或 os.Exit 前的显式调用。
• Helper 函数必须调用 t.Helper()：否则测试失败时错误行号指向 helper 内部，增加调试难度。
• 避免在 Helper 中调用 t.Fatal 和 t.Skip 之外的 t 方法：Helper 应专注于资源创建和清理，测试逻辑应留在测试函数中。
• 使用 t.Cleanup 而非 defer：t.Cleanup 在测试结束后按 LIFO 顺序执行，且能正确处理子测试的清理链。
• Kafka 资源命名隔离：在 TestMain 或 Helper 中为 Topic 和 Consumer Group 添加随机后缀（如 topic-{uuid}），避免并行测试时的命名冲突。