fix(postgres): Apply default snapshot name if no name specified (#2783)

testcontainers · Sep 17, 2024 · 060734b · 060734b
1 parent b60497e
commit 060734b
Show file tree

Hide file tree

Showing 3 changed files with 83 additions and 58 deletions.
diff --git a/modules/postgres/options.go b/modules/postgres/options.go
@@ -7,11 +7,13 @@ import (
 type options struct {
 	// SQLDriverName is the name of the SQL driver to use.
 	SQLDriverName string
+	Snapshot      string
 }
 
 func defaultOptions() options {
 	return options{
 		SQLDriverName: "postgres",
+		Snapshot:      defaultSnapshotName,
 	}
 }
 

diff --git a/modules/postgres/postgres.go b/modules/postgres/postgres.go
@@ -177,6 +177,7 @@ func Run(ctx context.Context, img string, opts ...testcontainers.ContainerCustom
 			password:      req.Env["POSTGRES_PASSWORD"],
 			user:          req.Env["POSTGRES_USER"],
 			sqlDriverName: settings.SQLDriverName,
+			snapshotName:  settings.Snapshot,
 		}
 	}
 

diff --git a/modules/postgres/postgres_test.go b/modules/postgres/postgres_test.go
@@ -203,73 +203,95 @@ func TestWithInitScript(t *testing.T) {
 }
 
 func TestSnapshot(t *testing.T) {
-	// snapshotAndReset {
-	ctx := context.Background()
-
-	// 1. Start the postgres ctr and run any migrations on it
-	ctr, err := postgres.Run(
-		ctx,
-		"docker.io/postgres:16-alpine",
-		postgres.WithDatabase(dbname),
-		postgres.WithUsername(user),
-		postgres.WithPassword(password),
-		postgres.BasicWaitStrategies(),
-		postgres.WithSQLDriver("pgx"),
-	)
-	testcontainers.CleanupContainer(t, ctr)
-	require.NoError(t, err)
-
-	// Run any migrations on the database
-	_, _, err = ctr.Exec(ctx, []string{"psql", "-U", user, "-d", dbname, "-c", "CREATE TABLE users (id SERIAL, name TEXT NOT NULL, age INT NOT NULL)"})
-	require.NoError(t, err)
+	tests := []struct {
+		name    string
+		options []postgres.SnapshotOption
+	}{
+		{
+			name:    "snapshot/default",
+			options: nil,
+		},
 
-	// 2. Create a snapshot of the database to restore later
-	err = ctr.Snapshot(ctx, postgres.WithSnapshotName("test-snapshot"))
-	require.NoError(t, err)
+		{
+			name: "snapshot/custom",
+			options: []postgres.SnapshotOption{
+				postgres.WithSnapshotName("custom-snapshot"),
+			},
+		},
+	}
 
-	dbURL, err := ctr.ConnectionString(ctx)
-	require.NoError(t, err)
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// snapshotAndReset {
+			ctx := context.Background()
 
-	t.Run("Test inserting a user", func(t *testing.T) {
-		t.Cleanup(func() {
-			// 3. In each test, reset the DB to its snapshot state.
-			err = ctr.Restore(ctx)
+			// 1. Start the postgres ctr and run any migrations on it
+			ctr, err := postgres.Run(
+				ctx,
+				"docker.io/postgres:16-alpine",
+				postgres.WithDatabase(dbname),
+				postgres.WithUsername(user),
+				postgres.WithPassword(password),
+				postgres.BasicWaitStrategies(),
+				postgres.WithSQLDriver("pgx"),
+			)
+			testcontainers.CleanupContainer(t, ctr)
 			require.NoError(t, err)
-		})
 
-		conn, err := pgx.Connect(context.Background(), dbURL)
-		require.NoError(t, err)
-		defer conn.Close(context.Background())
-
-		_, err = conn.Exec(ctx, "INSERT INTO users(name, age) VALUES ($1, $2)", "test", 42)
-		require.NoError(t, err)
-
-		var name string
-		var age int64
-		err = conn.QueryRow(context.Background(), "SELECT name, age FROM users LIMIT 1").Scan(&name, &age)
-		require.NoError(t, err)
-
-		require.Equal(t, "test", name)
-		require.EqualValues(t, 42, age)
-	})
+			// Run any migrations on the database
+			_, _, err = ctr.Exec(ctx, []string{"psql", "-U", user, "-d", dbname, "-c", "CREATE TABLE users (id SERIAL, name TEXT NOT NULL, age INT NOT NULL)"})
+			require.NoError(t, err)
 
-	// 4. Run as many tests as you need, they will each get a clean database
-	t.Run("Test querying empty DB", func(t *testing.T) {
-		t.Cleanup(func() {
-			err = ctr.Restore(ctx)
+			// 2. Create a snapshot of the database to restore later
+			// tt.options comes the test case, it can be specified as e.g. `postgres.WithSnapshotName("custom-snapshot")` or omitted, to use default name
+			err = ctr.Snapshot(ctx, tt.options...)
 			require.NoError(t, err)
-		})
 
-		conn, err := pgx.Connect(context.Background(), dbURL)
-		require.NoError(t, err)
-		defer conn.Close(context.Background())
+			dbURL, err := ctr.ConnectionString(ctx)
+			require.NoError(t, err)
 
-		var name string
-		var age int64
-		err = conn.QueryRow(context.Background(), "SELECT name, age FROM users LIMIT 1").Scan(&name, &age)
-		require.ErrorIs(t, err, pgx.ErrNoRows)
-	})
-	// }
+			t.Run("Test inserting a user", func(t *testing.T) {
+				t.Cleanup(func() {
+					// 3. In each test, reset the DB to its snapshot state.
+					err = ctr.Restore(ctx)
+					require.NoError(t, err)
+				})
+
+				conn, err := pgx.Connect(context.Background(), dbURL)
+				require.NoError(t, err)
+				defer conn.Close(context.Background())
+
+				_, err = conn.Exec(ctx, "INSERT INTO users(name, age) VALUES ($1, $2)", "test", 42)
+				require.NoError(t, err)
+
+				var name string
+				var age int64
+				err = conn.QueryRow(context.Background(), "SELECT name, age FROM users LIMIT 1").Scan(&name, &age)
+				require.NoError(t, err)
+
+				require.Equal(t, "test", name)
+				require.EqualValues(t, 42, age)
+			})
+
+			// 4. Run as many tests as you need, they will each get a clean database
+			t.Run("Test querying empty DB", func(t *testing.T) {
+				t.Cleanup(func() {
+					err = ctr.Restore(ctx)
+					require.NoError(t, err)
+				})
+
+				conn, err := pgx.Connect(context.Background(), dbURL)
+				require.NoError(t, err)
+				defer conn.Close(context.Background())
+
+				var name string
+				var age int64
+				err = conn.QueryRow(context.Background(), "SELECT name, age FROM users LIMIT 1").Scan(&name, &age)
+				require.ErrorIs(t, err, pgx.ErrNoRows)
+			})
+			// }
+		})
+	}
 }
 
 func TestSnapshotWithOverrides(t *testing.T) {