大多数 Go 测试只需要普通 TestXxx。但有时一个包里的多组测试需要共享准备工作,比如创建临时目录、启动测试服务器、准备测试数据库连接。TestMain 可以在一个测试包运行前后执行统一逻辑。
本文讲 TestMain 的基本结构和边界。它有用,但不要滥用。很多测试用 t.TempDir、helper 和子测试就够了。
最小 TestMain
func TestMain(m *testing.M) {
code := m.Run()
os.Exit(code)
}
m.Run() 会运行当前包的测试,返回退出码。你可以在它前面准备环境,在后面清理。最后必须 os.Exit(code),否则退出码可能不正确。
共享临时目录
var testDataDir string
func TestMain(m *testing.M) {
dir, err := os.MkdirTemp("", "myapp-test-*")
if err != nil {
fmt.Fprintln(os.Stderr, err)
os.Exit(1)
}
testDataDir = dir
code := m.Run()
os.RemoveAll(dir)
os.Exit(code)
}
测试里使用:
func TestReadFixture(t *testing.T) {
path := filepath.Join(testDataDir, "input.txt")
if err := os.WriteFile(path, []byte("hello"), 0644); err != nil {
t.Fatal(err)
}
}
如果每个测试都可以独立目录,优先用 t.TempDir()。共享目录适合准备成本高、内容只读的 fixture。共享可写状态容易让测试互相影响。
测试服务器
var testServerURL string
var closeServer func()
func TestMain(m *testing.M) {
server := httptest.NewServer(routes())
testServerURL = server.URL
closeServer = server.Close
code := m.Run()
closeServer()
os.Exit(code)
}
这种方式适合集成风格测试。普通 handler 测试仍然可以直接用 httptest.NewRecorder,不一定要启动 server。TestMain 准备越重,测试包运行越慢。
环境变量
TestMain 没有 *testing.T,所以不能用 t.Setenv。需要手动保存和恢复,或者在 m.Run 前设置:
old := os.Getenv("APP_ENV")
os.Setenv("APP_ENV", "test")
code := m.Run()
os.Setenv("APP_ENV", old)
os.Exit(code)
如果只影响单个测试,还是用 t.Setenv 更好。TestMain 设置的是整个包级环境,范围更大。
数据库测试
如果包里的测试都需要数据库,可以在 TestMain 里检查环境变量:
var testDB *sql.DB
func TestMain(m *testing.M) {
dsn := os.Getenv("TEST_DATABASE_URL")
if dsn == "" {
fmt.Fprintln(os.Stderr, "TEST_DATABASE_URL is required")
os.Exit(1)
}
db, err := sql.Open("postgres", dsn)
if err != nil {
fmt.Fprintln(os.Stderr, err)
os.Exit(1)
}
testDB = db
code := m.Run()
db.Close()
os.Exit(code)
}
这类测试更像集成测试,应该在文档里说明如何运行。不要让普通单元测试无意中依赖外部数据库。
不要把断言写进 TestMain
TestMain 主要负责 setup 和 teardown。具体行为断言仍然应该在 TestXxx 里。否则失败信息会很粗糙,也无法使用 t.Helper、子测试、临时目录等测试工具。
如果 setup 失败,只能打印到 stderr 并退出:
fmt.Fprintln(os.Stderr, "setup failed:", err)
os.Exit(1)
所以 TestMain 里逻辑要尽量少。
和并行测试的关系
TestMain 只负责整个包的入口和出口,并不会替你解决并行测试的数据隔离。如果测试里调用了 t.Parallel(),共享目录、共享数据库记录、共享环境变量都可能互相影响。
func TestWriteProfile(t *testing.T) {
t.Parallel()
dir := t.TempDir()
path := filepath.Join(dir, "profile.json")
if err := os.WriteFile(path, []byte(`{"name":"li"}`), 0644); err != nil {
t.Fatal(err)
}
}
即使有 TestMain,每个测试仍然应该尽量使用 t.TempDir()、唯一 ID 和独立数据。TestMain 适合准备共享的重资源,比如启动一次测试数据库、加载一次证书、初始化一个假服务地址,而不是把所有临时文件都放在同一个目录里。
清理失败也要处理
很多示例会在 TestMain 里简单写 defer cleanup(),但真实项目里清理也可能失败。比如删除临时目录失败、关闭测试服务失败、数据库容器停止失败。测试主体已经失败时,清理失败也要让日志可见。
func TestMain(m *testing.M) {
dir, err := os.MkdirTemp("", "demo-fixture-*")
if err != nil {
fmt.Fprintln(os.Stderr, err)
os.Exit(1)
}
code := m.Run()
if err := os.RemoveAll(dir); err != nil {
fmt.Fprintf(os.Stderr, "cleanup fixture dir: %v\n", err)
if code == 0 {
code = 1
}
}
os.Exit(code)
}
这个写法会保留原始测试结果,同时在清理失败且测试原本成功时把退出码改成失败。对 CI 来说,这比悄悄留下临时资源更可靠。
环境变量要恢复
如果测试需要改环境变量,优先在单个测试里用 t.Setenv。它会在测试结束时恢复旧值,比手动 os.Setenv 安全。
func TestConfigFromEnv(t *testing.T) {
t.Setenv("APP_MODE", "test")
cfg := LoadConfig()
if cfg.Mode != "test" {
t.Fatalf("mode = %q", cfg.Mode)
}
}
TestMain 里设置环境变量要更谨慎,因为它会影响整个包的所有测试。除非这是包级约定,比如统一把 APP_ENV 设成 test,否则不要在入口里藏太多状态。
夹具数据保持小而清楚
测试夹具不是越像生产越好。一个 3MB 的 JSON 文件也许能覆盖真实场景,但新人读测试时会很难理解失败原因。更好的方式是准备几份小数据:正常数据、缺字段数据、边界值数据、损坏数据。
var validUserJSON = []byte(`{"id":1,"name":"nina","active":true}`)
var missingNameJSON = []byte(`{"id":1,"active":true}`)
夹具的目标是解释行为,而不是复制生产库。只有在测试解析性能、兼容历史导出文件时,才需要引入较大的真实样本。
小结
TestMain 适合为一个测试包准备共享环境,比如测试服务器、临时目录、数据库连接。基本结构是 setup、m.Run()、teardown、os.Exit(code)。
不要滥用 TestMain。单个测试能用 t.TempDir、t.Setenv、helper 解决,就不要提升到包级共享。共享环境越多,测试之间越容易互相影响。
继续阅读
探索更多技术文章
浏览归档,发现更多关于系统设计、工具链和工程实践的内容。