🗂

DBアクセスがあるユニットテストをトランザクションを使って倍速にした

2023/12/21に公開

これは CastingONE Advent Calendar 2023 21日目の記事です。

株式会社CastingONEでバックエンドエンジニアをしている村上です。

弊社システムのDBをMySQLからPostgreSQLに移行した際、以前は数分で終わっていたCIがゆうに10分を超えるようになりました。実行時間の大部分はDBアクセスを伴うユニットテストで、テストの書き方を見直すことで劇的に時間短縮することができました。備忘録としてその方法をまとめます。

Before: テストケースごとにレコードをDELETE & INSERT

各テストの最初にDBの既存レコードをDELETE & テスト用レコードをINSERT(いわゆるsetUp & tearDown)していて、このテストデータを準備する部分に時間がかかっていました。

テストデータが増えると1回あたりの実行時間が増えますし、テストケースが増えるとそれが倍々に延びていきます。

当時の状況を再現します。

当時の状況

サンプルデータを入れるテーブルとして、usersテーブルを用意しました。

users.sql
CREATE TABLE users (
  id integer, 
  name varchar(100),
  memo text
);

テストデータをyamlで定義しています。
それなりの大きさのデータが1万件あります。

これを下記のライブラリを使ってINSERTします。
https://github.com/go-testfixtures/testfixtures

users.yml
- id: 0
  name: user_0
  memo: memo_0_T9cDouHaUOk37zfA1bFmthWS86kOHRA0QrTMHZg71YvH4ERW9t6x6a64k6J1XrZVFn6JnO5A1ecbA2GeXuh7zQpOLUZE8a489Icc1OUkNJYKdsIVD2GNrKZ0PoADNwyvuxHag6256vBTKWGANwARbVUn4k6W13r7xdz8Z4uhWRD98Cj8MiuF4V2TxKuOfuECdTj3F0bnu6T55bzNsggSk8svlb268dtWhIHJGgsK10b88ZpumBopwBPI1VaPRQkRHiD7KWUiOoVpoFUvM8TGZ3IBMHPd4IQWiHMExAarYSUb2TIDtXGEkcodUcWahwjXKn0Gjch5VTkOizwrhRU2ysCHFZVqHHfkMhCn4vJGWh5uVxIx7MzeZWgjSGDbGhC5Zx6Q1zrWamF7PN3t7jdGsYDNFxZu7A24xvyUksJAdV2NHev0oTVvvGz9AP4lPm71mnU3xD9UrywfHfaQNqpNIRAgtuKI08CuCeOkraTEYV2uL76CLodh3um1ipkc8gK1Td8PVRU97D0UsmWD5LxF1x2onhqd5OJVWxMFCNoagSJqulrtzgpkVGMfIEPN6sBxSCnq9CM4oLKuuz8TB9SvwGNQ0w5yAIQRY0iSgLwe1ykl5S28REnuvaqNDYgd5roOMs7SGbLYLpfRH1vYHwBE8X76YZ3X6273irHhgh2KR2rJI8WgzBc380CPGRfhw6jTpiBVazS5r00T3gaxG3kyO85SVtjccQBEeJjvzFkh9UyKtEKjw7OJyIMuAfzYcSSTS65sFZYa3j12uC0ObDN2NBmj1fWLY3LjeVdyofksO3JYOUWrW51rgvS3CslH70vN6U5vi6zr4s99KrDqHT9dyobn4rjPJffmkQCJRxjZKoDPn2j5tldyDG2cnshqAuwPtodpNbppKtGYNM0fIfYFSlaqojkxrHZK9iygcppIKv8V9U2cm1NU3jSANpNIKQr51WGvJnYRTGGoVwHS1tgwktuRWqaEvMamlYmgygJF
- id: 1
  name: user_1
  memo: memo_1_T9cDouHaUOk37zfA1bFmthWS86kOHRA0QrTMHZg71YvH4ERW9t6x6a64k6J1XrZVFn6JnO5A1ecbA2GeXuh7zQpOLUZE8a489Icc1OUkNJYKdsIVD2GNrKZ0PoADNwyvuxHag6256vBTKWGANwARbVUn4k6W13r7xdz8Z4uhWRD98Cj8MiuF4V2TxKuOfuECdTj3F0bnu6T55bzNsggSk8svlb268dtWhIHJGgsK10b88ZpumBopwBPI1VaPRQkRHiD7KWUiOoVpoFUvM8TGZ3IBMHPd4IQWiHMExAarYSUb2TIDtXGEkcodUcWahwjXKn0Gjch5VTkOizwrhRU2ysCHFZVqHHfkMhCn4vJGWh5uVxIx7MzeZWgjSGDbGhC5Zx6Q1zrWamF7PN3t7jdGsYDNFxZu7A24xvyUksJAdV2NHev0oTVvvGz9AP4lPm71mnU3xD9UrywfHfaQNqpNIRAgtuKI08CuCeOkraTEYV2uL76CLodh3um1ipkc8gK1Td8PVRU97D0UsmWD5LxF1x2onhqd5OJVWxMFCNoagSJqulrtzgpkVGMfIEPN6sBxSCnq9CM4oLKuuz8TB9SvwGNQ0w5yAIQRY0iSgLwe1ykl5S28REnuvaqNDYgd5roOMs7SGbLYLpfRH1vYHwBE8X76YZ3X6273irHhgh2KR2rJI8WgzBc380CPGRfhw6jTpiBVazS5r00T3gaxG3kyO85SVtjccQBEeJjvzFkh9UyKtEKjw7OJyIMuAfzYcSSTS65sFZYa3j12uC0ObDN2NBmj1fWLY3LjeVdyofksO3JYOUWrW51rgvS3CslH70vN6U5vi6zr4s99KrDqHT9dyobn4rjPJffmkQCJRxjZKoDPn2j5tldyDG2cnshqAuwPtodpNbppKtGYNM0fIfYFSlaqojkxrHZK9iygcppIKv8V9U2cm1NU3jSANpNIKQr51WGvJnYRTGGoVwHS1tgwktuRWqaEvMamlYmgygJF

  # id: 2~9997 は省略

- id: 9998
  name: user_9998
  memo: memo_9998_T9cDouHaUOk37zfA1bFmthWS86kOHRA0QrTMHZg71YvH4ERW9t6x6a64k6J1XrZVFn6JnO5A1ecbA2GeXuh7zQpOLUZE8a489Icc1OUkNJYKdsIVD2GNrKZ0PoADNwyvuxHag6256vBTKWGANwARbVUn4k6W13r7xdz8Z4uhWRD98Cj8MiuF4V2TxKuOfuECdTj3F0bnu6T55bzNsggSk8svlb268dtWhIHJGgsK10b88ZpumBopwBPI1VaPRQkRHiD7KWUiOoVpoFUvM8TGZ3IBMHPd4IQWiHMExAarYSUb2TIDtXGEkcodUcWahwjXKn0Gjch5VTkOizwrhRU2ysCHFZVqHHfkMhCn4vJGWh5uVxIx7MzeZWgjSGDbGhC5Zx6Q1zrWamF7PN3t7jdGsYDNFxZu7A24xvyUksJAdV2NHev0oTVvvGz9AP4lPm71mnU3xD9UrywfHfaQNqpNIRAgtuKI08CuCeOkraTEYV2uL76CLodh3um1ipkc8gK1Td8PVRU97D0UsmWD5LxF1x2onhqd5OJVWxMFCNoagSJqulrtzgpkVGMfIEPN6sBxSCnq9CM4oLKuuz8TB9SvwGNQ0w5yAIQRY0iSgLwe1ykl5S28REnuvaqNDYgd5roOMs7SGbLYLpfRH1vYHwBE8X76YZ3X6273irHhgh2KR2rJI8WgzBc380CPGRfhw6jTpiBVazS5r00T3gaxG3kyO85SVtjccQBEeJjvzFkh9UyKtEKjw7OJyIMuAfzYcSSTS65sFZYa3j12uC0ObDN2NBmj1fWLY3LjeVdyofksO3JYOUWrW51rgvS3CslH70vN6U5vi6zr4s99KrDqHT9dyobn4rjPJffmkQCJRxjZKoDPn2j5tldyDG2cnshqAuwPtodpNbppKtGYNM0fIfYFSlaqojkxrHZK9iygcppIKv8V9U2cm1NU3jSANpNIKQr51WGvJnYRTGGoVwHS1tgwktuRWqaEvMamlYmgygJF
- id: 9999
  name: user_9999
  memo: memo_9999_T9cDouHaUOk37zfA1bFmthWS86kOHRA0QrTMHZg71YvH4ERW9t6x6a64k6J1XrZVFn6JnO5A1ecbA2GeXuh7zQpOLUZE8a489Icc1OUkNJYKdsIVD2GNrKZ0PoADNwyvuxHag6256vBTKWGANwARbVUn4k6W13r7xdz8Z4uhWRD98Cj8MiuF4V2TxKuOfuECdTj3F0bnu6T55bzNsggSk8svlb268dtWhIHJGgsK10b88ZpumBopwBPI1VaPRQkRHiD7KWUiOoVpoFUvM8TGZ3IBMHPd4IQWiHMExAarYSUb2TIDtXGEkcodUcWahwjXKn0Gjch5VTkOizwrhRU2ysCHFZVqHHfkMhCn4vJGWh5uVxIx7MzeZWgjSGDbGhC5Zx6Q1zrWamF7PN3t7jdGsYDNFxZu7A24xvyUksJAdV2NHev0oTVvvGz9AP4lPm71mnU3xD9UrywfHfaQNqpNIRAgtuKI08CuCeOkraTEYV2uL76CLodh3um1ipkc8gK1Td8PVRU97D0UsmWD5LxF1x2onhqd5OJVWxMFCNoagSJqulrtzgpkVGMfIEPN6sBxSCnq9CM4oLKuuz8TB9SvwGNQ0w5yAIQRY0iSgLwe1ykl5S28REnuvaqNDYgd5roOMs7SGbLYLpfRH1vYHwBE8X76YZ3X6273irHhgh2KR2rJI8WgzBc380CPGRfhw6jTpiBVazS5r00T3gaxG3kyO85SVtjccQBEeJjvzFkh9UyKtEKjw7OJyIMuAfzYcSSTS65sFZYa3j12uC0ObDN2NBmj1fWLY3LjeVdyofksO3JYOUWrW51rgvS3CslH70vN6U5vi6zr4s99KrDqHT9dyobn4rjPJffmkQCJRxjZKoDPn2j5tldyDG2cnshqAuwPtodpNbppKtGYNM0fIfYFSlaqojkxrHZK9iygcppIKv8V9U2cm1NU3jSANpNIKQr51WGvJnYRTGGoVwHS1tgwktuRWqaEvMamlYmgygJF

テストケースが4つ(TestInsert, TestUpdate, TestDelete, TestSelect)あり、それぞれ最初にprepareTestDatabaseを呼ぶことで上記のyamlをusersテーブルにINSERTしています。TestMainはテストデータを入れるLoaderを用意しているだけで、その場でINSERTはしていません。

main_test.go
package main

import (
	"database/sql"
	"log"
	"os"
	"testing"

	"github.com/go-testfixtures/testfixtures/v3"
	_ "github.com/lib/pq"
)

var (
	// FixtureをDBにロードするためのLoader
	loader *testfixtures.Loader

	db *sql.DB
)

func TestMain(m *testing.M) {
	// DB接続
	dsn := "user=postgres password=postgres database=postgres host=127.0.0.1 search_path=test_schema sslmode=disable"
	var err error
	db, err = sql.Open("postgres", dsn)
	if err != nil {
		log.Fatal(err)
	}

	// Loaderの準備
	loader, err = testfixtures.New(
		testfixtures.Database(db),
		testfixtures.Dialect("postgres"),
		testfixtures.Files("users.yml"),
		testfixtures.DangerousSkipTestDatabaseCheck(),
	)
	if err != nil {
		log.Fatal(err)
	}

	os.Exit(m.Run())
}

// FixtureをDBにロードする
func prepareTestDatabase() {
	if err := loader.Load(); err != nil {
		log.Fatal(err)
	}
}

// usersのレコードを1件追加し、レコード件数がFixtureの件数+1になることを確認する
func TestInsert(t *testing.T) {
	prepareTestDatabase()

	if _, err := db.Exec("INSERT INTO users (name) VALUES ('test')"); err != nil {
		t.Fatal(err)
	}

	var count int
	if err := db.QueryRow("SELECT COUNT(*) FROM users").Scan(&count); err != nil {
		t.Fatal(err)
	}
	if count != 10001 {
		t.Errorf("expected count to be 10001, but got %d", count)
	}
}

// usersのレコードを1件更新し、レコード件数がFixtureの件数と変わらないことを確認する
func TestUpdate(t *testing.T) {
	prepareTestDatabase()

	if _, err := db.Exec("UPDATE users SET name = 'test' WHERE id = 9999"); err != nil {
		t.Fatal(err)
	}

	var count int
	if err := db.QueryRow("SELECT COUNT(*) FROM users").Scan(&count); err != nil {
		t.Fatal(err)
	}
	if count != 10000 {
		t.Errorf("expected count to be 10000, but got %d", count)
	}
}

// usersのレコードを1件削除し、レコード件数がFixtureの件数-1になることを確認する
func TestDelete(t *testing.T) {
	prepareTestDatabase()

	if _, err := db.Exec("DELETE FROM users WHERE id = 9999"); err != nil {
		t.Fatal(err)
	}

	var count int
	if err := db.QueryRow("SELECT COUNT(*) FROM users").Scan(&count); err != nil {
		t.Fatal(err)
	}
	if count != 9999 {
		t.Errorf("expected count to be 9999, but got %d", count)
	}
}

// usersのレコード件数がFixtureの件数と一致することを確認する
func TestSelect(t *testing.T) {
	prepareTestDatabase()

	var count int
	if err := db.QueryRow("SELECT COUNT(*) FROM users").Scan(&count); err != nil {
		t.Fatal(err)
	}
	if count != 10000 {
		t.Errorf("expected count to be 10000, but got %d", count)
	}
}

テストを実行すると、テストケースごとに14秒前後かかっていることが分かります。テストケースが10個になると140秒、100個になると1400秒...と増えていくので、大変です。

$ go test -v
=== RUN   TestInsert
--- PASS: TestInsert (14.82s)
=== RUN   TestUpdate
--- PASS: TestUpdate (14.24s)
=== RUN   TestDelete
--- PASS: TestDelete (13.89s)
=== RUN   TestSelect
--- PASS: TestSelect (13.88s)
PASS
ok      rollback_testdata       57.511s

After: テスト後にロールバック

テストケースごとにDELETE & INSERTすることに時間がかかっていたため、テストを開始する前にテストデータをINSERTし、各テストで生じたデータの差分はロールバックするようにしました。これにより、DBのI/Oを抑えた上で、今までと同じようにテストデータを初期化できるようになりました。

具体的な変更点は、下記のとおりです。

  • TestMainでテストデータをINSERTするようにした
  • TestInsert, TestUpdate, TestDelete, TestSelectでテストデータをINSERTしないようにした。替わりに、テスト終了時にロールバックするようにした。
package main

import (
	"database/sql"
	"log"
	"os"
	"testing"

	"github.com/go-testfixtures/testfixtures/v3"
	_ "github.com/lib/pq"
)

var (
	// FixtureをDBにロードするためのLoader
	loader *testfixtures.Loader

	db *sql.DB
)

func TestMain(m *testing.M) {
	// DB接続
	dsn := "user=postgres password=postgres database=postgres host=127.0.0.1 search_path=test_schema sslmode=disable"
	var err error
	db, err = sql.Open("postgres", dsn)
	if err != nil {
		log.Fatal(err)
	}

	// Loaderの準備
	loader, err = testfixtures.New(
		testfixtures.Database(db),
		testfixtures.Dialect("postgres"),
		testfixtures.Files("users.yml"),
		testfixtures.DangerousSkipTestDatabaseCheck(),
	)
	if err != nil {
		log.Fatal(err)
	}

	// テストの開始前に1回だけFixtureをロードする
	prepareTestDatabase()

	os.Exit(m.Run())
}

// FixtureをDBにロードする
func prepareTestDatabase() {
	if err := loader.Load(); err != nil {
		log.Fatal(err)
	}
}

// usersのレコードを1件追加し、レコード件数がFixtureの件数+1になることを確認する
func TestInsert(t *testing.T) {
	// テストで生じたDBの変更をロールバックする
	tx, err := db.Begin()
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		tx.Rollback()
	})

	// 個々のテストではFixtureをロードする必要がなくなった
	// prepareTestDatabase()

	if _, err := tx.Exec("INSERT INTO users (name) VALUES ('test')"); err != nil {
		t.Fatal(err)
	}

	var count int
	if err := tx.QueryRow("SELECT COUNT(*) FROM users").Scan(&count); err != nil {
		t.Fatal(err)
	}
	if count != 10001 {
		t.Errorf("expected count to be 10001, but got %d", count)
	}
}

// usersのレコードを1件更新し、レコード件数がFixtureの件数と変わらないことを確認する
func TestUpdate(t *testing.T) {
	// テストで生じたDBの変更をロールバックする
	tx, err := db.Begin()
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		tx.Rollback()
	})

	// 個々のテストではFixtureをロードする必要がなくなった
	// prepareTestDatabase()

	if _, err := tx.Exec("UPDATE users SET name = 'test' WHERE id = 9999"); err != nil {
		t.Fatal(err)
	}

	var count int
	if err := tx.QueryRow("SELECT COUNT(*) FROM users").Scan(&count); err != nil {
		t.Fatal(err)
	}
	if count != 10000 {
		t.Errorf("expected count to be 10000, but got %d", count)
	}
}

// usersのレコードを1件削除し、レコード件数がFixtureの件数-1になることを確認する
func TestDelete(t *testing.T) {
	// テストで生じたDBの変更をロールバックする
	tx, err := db.Begin()
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		tx.Rollback()
	})

	// 個々のテストではFixtureをロードする必要がなくなった
	// prepareTestDatabase()

	if _, err := tx.Exec("DELETE FROM users WHERE id = 9999"); err != nil {
		t.Fatal(err)
	}

	var count int
	if err := tx.QueryRow("SELECT COUNT(*) FROM users").Scan(&count); err != nil {
		t.Fatal(err)
	}
	if count != 9999 {
		t.Errorf("expected count to be 9999, but got %d", count)
	}
}

// usersのレコード件数がFixtureの件数と一致することを確認する
func TestSelect(t *testing.T) {
	// テストで生じたDBの変更をロールバックする
	tx, err := db.Begin()
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		tx.Rollback()
	})

	// 個々のテストではFixtureをロードする必要がなくなった
	// prepareTestDatabase()

	var count int
	if err := tx.QueryRow("SELECT COUNT(*) FROM users").Scan(&count); err != nil {
		t.Fatal(err)
	}
	if count != 10000 {
		t.Errorf("expected count to be 10000, but got %d", count)
	}
}

テストを実行すると、最初の1回だけテストデータのINSERTに14秒を要していますが、各テストケースは0秒まで下がる結果となりました。

$ go test -v
=== RUN   TestInsert
--- PASS: TestInsert (0.00s)
=== RUN   TestUpdate
--- PASS: TestUpdate (0.00s)
=== RUN   TestDelete
--- PASS: TestDelete (0.00s)
=== RUN   TestSelect
--- PASS: TestSelect (0.00s)
PASS
ok      rollback_testdata       14.829s

終わりに

テストごとにデータを入れ直すことをやめて、ロールバックすることでデータを初期化するようにしたことで、テストの実行時間をグッと抑えることができました。今回のサンプルコードでは標準packageを使いましたが、AfterEachなどを使えるライブラリと組み合わせれば、ロールバックの部分も簡潔に書けるはずです。何かの参考になれば嬉しいです。

弊社ではいっしょに働いてくれるエンジニアを募集中です。社員でもフリーランスでも、フルタイムでも短時間の副業でも大歓迎なので、気軽にご連絡ください!

https://www.wantedly.com/projects/1130967

Discussion