From dbb9c71df8d3679f3f25cb229fd104862123963a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@blender.org>
Date: Fri, 4 Mar 2022 12:33:45 +0100
Subject: [PATCH] Tests: more unified way to do database tests

---
 internal/manager/persistence/jobs_test.go     | 12 ++-----
 .../persistence/task_scheduler_test.go        | 32 +++++++------------
 internal/manager/persistence/test_support.go  | 16 ++++++++++
 internal/manager/persistence/workers_test.go  |  9 ++----
 4 files changed, 31 insertions(+), 38 deletions(-)

diff --git a/internal/manager/persistence/jobs_test.go b/internal/manager/persistence/jobs_test.go
index d94a0ff7..d7f153a8 100644
--- a/internal/manager/persistence/jobs_test.go
+++ b/internal/manager/persistence/jobs_test.go
@@ -33,10 +33,7 @@ import (
 )
 
 func TestStoreAuthoredJob(t *testing.T) {
-	db, dbCloser := CreateTestDB(t)
-	defer dbCloser()
-
-	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	ctx, cancel, db := persistenceTestFixtures(t, 1*time.Second)
 	defer cancel()
 
 	job := createTestAuthoredJobWithTasks()
@@ -297,12 +294,7 @@ func createTestAuthoredJobWithTasks() job_compilers.AuthoredJob {
 }
 
 func jobTasksTestFixtures(t *testing.T) (context.Context, context.CancelFunc, *DB, *Job, job_compilers.AuthoredJob) {
-	db, dbCloser := CreateTestDB(t)
-	ctx, ctxCancel := context.WithTimeout(context.Background(), 1*time.Second)
-	cancel := func() {
-		ctxCancel()
-		dbCloser()
-	}
+	ctx, cancel, db := persistenceTestFixtures(t, schedulerTestTimeout)
 
 	authoredJob := createTestAuthoredJobWithTasks()
 	err := db.StoreAuthoredJob(ctx, authoredJob)
diff --git a/internal/manager/persistence/task_scheduler_test.go b/internal/manager/persistence/task_scheduler_test.go
index 57ef4bdb..3546ffc2 100644
--- a/internal/manager/persistence/task_scheduler_test.go
+++ b/internal/manager/persistence/task_scheduler_test.go
@@ -32,13 +32,11 @@ import (
 	"git.blender.org/flamenco/pkg/api"
 )
 
-const testContextTimeout = 100 * time.Millisecond
+const schedulerTestTimeout = 100 * time.Millisecond
 
 func TestNoTasks(t *testing.T) {
-	db, dbCloser := CreateTestDB(t)
-	defer dbCloser()
-	ctx, ctxCancel := context.WithTimeout(context.Background(), testContextTimeout)
-	defer ctxCancel()
+	ctx, cancel, db := persistenceTestFixtures(t, schedulerTestTimeout)
+	defer cancel()
 
 	w := linuxWorker(t, db)
 
@@ -48,10 +46,8 @@ func TestNoTasks(t *testing.T) {
 }
 
 func TestOneJobOneTask(t *testing.T) {
-	db, dbCloser := CreateTestDB(t)
-	defer dbCloser()
-	ctx, ctxCancel := context.WithTimeout(context.Background(), testContextTimeout)
-	defer ctxCancel()
+	ctx, cancel, db := persistenceTestFixtures(t, schedulerTestTimeout)
+	defer cancel()
 
 	w := linuxWorker(t, db)
 
@@ -85,10 +81,8 @@ func TestOneJobOneTask(t *testing.T) {
 }
 
 func TestOneJobThreeTasksByPrio(t *testing.T) {
-	db, dbCloser := CreateTestDB(t)
-	defer dbCloser()
-	ctx, ctxCancel := context.WithTimeout(context.Background(), testContextTimeout)
-	defer ctxCancel()
+	ctx, cancel, db := persistenceTestFixtures(t, schedulerTestTimeout)
+	defer cancel()
 
 	w := linuxWorker(t, db)
 
@@ -118,10 +112,8 @@ func TestOneJobThreeTasksByPrio(t *testing.T) {
 }
 
 func TestOneJobThreeTasksByDependencies(t *testing.T) {
-	db, dbCloser := CreateTestDB(t)
-	defer dbCloser()
-	ctx, ctxCancel := context.WithTimeout(context.Background(), testContextTimeout)
-	defer ctxCancel()
+	ctx, cancel, db := persistenceTestFixtures(t, schedulerTestTimeout)
+	defer cancel()
 
 	w := linuxWorker(t, db)
 
@@ -146,10 +138,8 @@ func TestOneJobThreeTasksByDependencies(t *testing.T) {
 }
 
 func TestTwoJobsThreeTasks(t *testing.T) {
-	db, dbCloser := CreateTestDB(t)
-	defer dbCloser()
-	ctx, ctxCancel := context.WithTimeout(context.Background(), testContextTimeout)
-	defer ctxCancel()
+	ctx, cancel, db := persistenceTestFixtures(t, schedulerTestTimeout)
+	defer cancel()
 
 	w := linuxWorker(t, db)
 
diff --git a/internal/manager/persistence/test_support.go b/internal/manager/persistence/test_support.go
index 67ce9ac8..6b6c1bae 100644
--- a/internal/manager/persistence/test_support.go
+++ b/internal/manager/persistence/test_support.go
@@ -22,9 +22,11 @@ package persistence
  * ***** END GPL LICENSE BLOCK ***** */
 
 import (
+	"context"
 	"database/sql"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/glebarez/sqlite"
 	"github.com/rs/zerolog"
@@ -78,3 +80,17 @@ func CreateTestDB(t *testing.T) (db *DB, closer func()) {
 
 	return db, closer
 }
+
+// persistenceTestFixtures creates a test database and returns it and a context.
+// Tests should call the returned cancel function when they're done.
+func persistenceTestFixtures(t *testing.T, testContextTimeout time.Duration) (context.Context, context.CancelFunc, *DB) {
+	db, dbCloser := CreateTestDB(t)
+	ctx, ctxCancel := context.WithTimeout(context.Background(), testContextTimeout)
+
+	cancel := func() {
+		ctxCancel()
+		dbCloser()
+	}
+
+	return ctx, cancel, db
+}
diff --git a/internal/manager/persistence/workers_test.go b/internal/manager/persistence/workers_test.go
index 6d2da582..7f80a26c 100644
--- a/internal/manager/persistence/workers_test.go
+++ b/internal/manager/persistence/workers_test.go
@@ -27,15 +27,12 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
-	"golang.org/x/net/context"
 
 	"git.blender.org/flamenco/pkg/api"
 )
 
 func TestCreateFetchWorker(t *testing.T) {
-	db, dbCloser := CreateTestDB(t)
-	defer dbCloser()
-	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	ctx, cancel, db := persistenceTestFixtures(t, 1*time.Second)
 	defer cancel()
 
 	w := Worker{
@@ -69,9 +66,7 @@ func TestCreateFetchWorker(t *testing.T) {
 }
 
 func TestSaveWorker(t *testing.T) {
-	db, dbCloser := CreateTestDB(t)
-	defer dbCloser()
-	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	ctx, cancel, db := persistenceTestFixtures(t, 1*time.Second)
 	defer cancel()
 
 	w := Worker{