3 files changed, 6 insertions, 2 deletions
diff --git a/internal/postgres/insert_module_test.go b/internal/postgres/insert_module_test.go
index 955011a3..9807091e 100644
--- a/internal/postgres/insert_module_test.go
+++ b/internal/postgres/insert_module_test.go
@@ -22,6 +22,7 @@ import (
 	"golang.org/x/pkgsite/internal"
 	"golang.org/x/pkgsite/internal/database"
 	"golang.org/x/pkgsite/internal/derrors"
+	"golang.org/x/pkgsite/internal/experiment"
 	"golang.org/x/pkgsite/internal/licenses"
 	"golang.org/x/pkgsite/internal/source"
 	"golang.org/x/pkgsite/internal/stdlib"
@@ -31,6 +32,7 @@ import (
 func TestInsertModule(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), testTimeout*2)
 	defer cancel()
+	ctx = experiment.NewContext(ctx, internal.ExperimentInsertPackageSource)
 
 	for _, test := range []struct {
 		name   string
@@ -169,6 +171,7 @@ func TestInsertModuleLicenseCheck(t *testing.T) {
 func TestUpsertModule(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
 	defer cancel()
+	ctx = experiment.NewContext(ctx, internal.ExperimentInsertPackageSource)
 	m := sample.Module("upsert.org", "v1.2.3", "dir/p")
 
 	// Insert the module.
diff --git a/internal/postgres/search_test.go b/internal/postgres/search_test.go
index e5a82185..b5615664 100644
--- a/internal/postgres/search_test.go
+++ b/internal/postgres/search_test.go
@@ -294,7 +294,7 @@ func TestSearch(t *testing.T) {
 	for _, test := range tests {
 		t.Run(test.label, func(t *testing.T) {
 			defer ResetTestDB(testDB, t)
-			ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+			ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
 			defer cancel()
 			for _, m := range test.modules {
 				if err := testDB.InsertModule(ctx, m); err != nil {
diff --git a/internal/postgres/unit_test.go b/internal/postgres/unit_test.go
index 53603643..b0096353 100644
--- a/internal/postgres/unit_test.go
+++ b/internal/postgres/unit_test.go
@@ -23,7 +23,7 @@ import (
 func TestGetUnit(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
 	defer cancel()
-
+	ctx = experiment.NewContext(ctx, internal.ExperimentInsertPackageSource)
 	defer ResetTestDB(testDB, t)
 	InsertSampleDirectoryTree(ctx, t, testDB)
 
@@ -208,6 +208,7 @@ func TestGetUnitFieldSet(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
 	defer cancel()
 
+	ctx = experiment.NewContext(ctx, internal.ExperimentInsertPackageSource)
 	defer ResetTestDB(testDB, t)
 
 	// Add a module that has READMEs in a directory and a package.