Go TestMain 入门:为一组测试准备共享环境

讲 TestMain 的基本用法,用临时目录和测试数据库示例说明共享 fixture、退出码和不该滥用的边界。

大多数 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.TempDirt.Setenv、helper 解决,就不要提升到包级共享。共享环境越多,测试之间越容易互相影响。

继续阅读

探索更多技术文章

浏览归档,发现更多关于系统设计、工具链和工程实践的内容。

全部文章 返回首页