Adding database connection cleanup procedure on Resource Manager tests

server/config/configresource/config_resource_test.go

@@ -2,6 +2,7 @@ package configresource_test
 
 import (
 	"context"
+	"database/sql"
 	"os"
 	"testing"
 
@@ -100,8 +101,7 @@ func TestConfigResource(t *testing.T) {
 	rmtests.TestResourceType(t, rmtests.ResourceTypeTest{
 		ResourceTypeSingular: configresource.ResourceName,
 		ResourceTypePlural:   configresource.ResourceNamePlural,
-		RegisterManagerFn: func(router *mux.Router) resourcemanager.Manager {
-			db := testmock.CreateMigratedDatabase()
+		RegisterManagerFn: func(router *mux.Router, db *sql.DB) resourcemanager.Manager {
 			configRepo := configresource.NewRepository(db)
 
 			manager := resourcemanager.New[configresource.Config](

server/config/demoresource/demo_resource_test.go

@@ -2,14 +2,14 @@ package demoresource_test
 
 import (
 	"context"
+	"database/sql"
 	"testing"
 
 	"github.com/gorilla/mux"
 	"github.com/kubeshop/tracetest/server/config/demoresource"
 	"github.com/kubeshop/tracetest/server/pkg/id"
 	"github.com/kubeshop/tracetest/server/resourcemanager"
 	rmtests "github.com/kubeshop/tracetest/server/resourcemanager/testutil"
-	"github.com/kubeshop/tracetest/server/testmock"
 )
 
 func TestPokeshopDemoResource(t *testing.T) {
@@ -49,8 +49,7 @@ func TestPokeshopDemoResource(t *testing.T) {
 	rmtests.TestResourceType(t, rmtests.ResourceTypeTest{
 		ResourceTypeSingular: demoresource.ResourceName,
 		ResourceTypePlural:   demoresource.ResourceNamePlural,
-		RegisterManagerFn: func(router *mux.Router) resourcemanager.Manager {
-			db := testmock.CreateMigratedDatabase()
+		RegisterManagerFn: func(router *mux.Router, db *sql.DB) resourcemanager.Manager {
 			demoRepository := demoresource.NewRepository(db)
 
 			manager := resourcemanager.New[demoresource.Demo](
@@ -149,8 +148,7 @@ func TestOpenTelemetryStoreDemoResource(t *testing.T) {
 	rmtests.TestResourceType(t, rmtests.ResourceTypeTest{
 		ResourceTypeSingular: demoresource.ResourceName,
 		ResourceTypePlural:   demoresource.ResourceNamePlural,
-		RegisterManagerFn: func(router *mux.Router) resourcemanager.Manager {
-			db := testmock.CreateMigratedDatabase()
+		RegisterManagerFn: func(router *mux.Router, db *sql.DB) resourcemanager.Manager {
 			demoRepository := demoresource.NewRepository(db)
 
 			manager := resourcemanager.New[demoresource.Demo](

server/executor/pollingprofile/polling_profile_resource_test.go

@@ -1,22 +1,21 @@
 package pollingprofile_test
 
 import (
+	"database/sql"
 	"testing"
 
 	"github.com/gorilla/mux"
 	"github.com/kubeshop/tracetest/server/executor/pollingprofile"
 	"github.com/kubeshop/tracetest/server/pkg/id"
 	"github.com/kubeshop/tracetest/server/resourcemanager"
 	rmtests "github.com/kubeshop/tracetest/server/resourcemanager/testutil"
-	"github.com/kubeshop/tracetest/server/testmock"
 )
 
 func TestPollingProfileResource(t *testing.T) {
 	rmtests.TestResourceType(t, rmtests.ResourceTypeTest{
 		ResourceTypeSingular: pollingprofile.ResourceName,
 		ResourceTypePlural:   pollingprofile.ResourceNamePlural,
-		RegisterManagerFn: func(router *mux.Router) resourcemanager.Manager {
-			db := testmock.CreateMigratedDatabase()
+		RegisterManagerFn: func(router *mux.Router, db *sql.DB) resourcemanager.Manager {
 			pollingProfileRepo := pollingprofile.NewRepository(db)
 
 			manager := resourcemanager.New[pollingprofile.PollingProfile](

server/resourcemanager/main_test.go

@@ -0,0 +1,18 @@
+package resourcemanager_test
+
+import (
+	"os"
+	"testing"
+
+	"github.com/kubeshop/tracetest/server/testmock"
+)
+
+func TestMain(m *testing.M) {
+	testmock.StartTestEnvironment()
+
+	exitVal := m.Run()
+
+	testmock.StopTestEnvironment()
+
+	os.Exit(exitVal)
+}

server/resourcemanager/resource_manager_test.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/gorilla/mux"
 	"github.com/kubeshop/tracetest/server/pkg/id"
+	"github.com/kubeshop/tracetest/server/resourcemanager"
 	rm "github.com/kubeshop/tracetest/server/resourcemanager"
 	rmtests "github.com/kubeshop/tracetest/server/resourcemanager/testutil"
 	"github.com/stretchr/testify/mock"
@@ -74,7 +75,7 @@ func TestSampleResource(t *testing.T) {
 	rmtests.TestResourceTypeWithErrorOperations(t, rmtests.ResourceTypeTest{
 		ResourceTypeSingular: "SampleResource",
 		ResourceTypePlural:   "SampleResources",
-		RegisterManagerFn: func(router *mux.Router) rm.Manager {
+		RegisterManagerFn: func(router *mux.Router, db *sql.DB) resourcemanager.Manager {
 			mockManager := new(sampleResourceManager)
 			manager := rm.New[sampleResource](
 				"SampleResource",
@@ -223,7 +224,7 @@ func TestRestrictedResource(t *testing.T) {
 	rmtests.TestResourceTypeWithErrorOperations(t, rmtests.ResourceTypeTest{
 		ResourceTypeSingular: "RestrictedResource",
 		ResourceTypePlural:   "RestrictedResources",
-		RegisterManagerFn: func(router *mux.Router) rm.Manager {
+		RegisterManagerFn: func(router *mux.Router, db *sql.DB) rm.Manager {
 			mockManager := new(restrictedResourceManager)
 			manager := rm.New[sampleResource](
 				"RestrictedResource",

server/resourcemanager/testutil/test_provisioning.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/gorilla/mux"
 	"github.com/kubeshop/tracetest/server/resourcemanager"
+	"github.com/kubeshop/tracetest/server/testmock"
 	"github.com/stretchr/testify/require"
 	"gopkg.in/yaml.v2"
 )
@@ -19,7 +20,10 @@ const (
 func testProvisioning(t *testing.T, rt ResourceTypeTest) {
 	t.Run("Provisioning", func(t *testing.T) {
 		t.Run("Success", func(t *testing.T) {
-			manager := rt.RegisterManagerFn(mux.NewRouter())
+			db := testmock.CreateMigratedDatabase()
+			defer db.Close()
+
+			manager := rt.RegisterManagerFn(mux.NewRouter(), db)
 			if rt.Prepare != nil {
 				rt.Prepare(t, OperationProvisioningSuccess, manager)
 			}
@@ -34,7 +38,10 @@ func testProvisioning(t *testing.T, rt ResourceTypeTest) {
 		})
 
 		t.Run("UnacceptableType", func(t *testing.T) {
-			manager := rt.RegisterManagerFn(mux.NewRouter())
+			db := testmock.CreateMigratedDatabase()
+			defer db.Close()
+
+			manager := rt.RegisterManagerFn(mux.NewRouter(), db)
 			if rt.Prepare != nil {
 				rt.Prepare(t, OperationProvisioningTypeNotSupported, manager)
 			}

server/resourcemanager/testutil/test_resource.go

@@ -1,6 +1,7 @@
 package testutil
 
 import (
+	"database/sql"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -10,12 +11,13 @@ import (
 	"gotest.tools/v3/assert"
 
 	rm "github.com/kubeshop/tracetest/server/resourcemanager"
+	"github.com/kubeshop/tracetest/server/testmock"
 )
 
 type ResourceTypeTest struct {
 	ResourceTypeSingular string
 	ResourceTypePlural   string
-	RegisterManagerFn    func(*mux.Router) rm.Manager
+	RegisterManagerFn    func(*mux.Router, *sql.DB) rm.Manager
 	Prepare              func(t *testing.T, operation Operation, manager rm.Manager)
 
 	SampleJSON        string
@@ -119,9 +121,15 @@ func testOperation(t *testing.T, op operationTester, rt ResourceTypeTest) {
 func testOperationForContentType(t *testing.T, op operationTester, ct contentTypeConverter, rt ResourceTypeTest) {
 	t.Helper()
 
+	db := testmock.CreateMigratedDatabase()
+	defer db.Close()
+
 	router := mux.NewRouter()
+
 	testServer := httptest.NewServer(router)
-	manager := rt.RegisterManagerFn(router)
+	defer testServer.Close()
+
+	manager := rt.RegisterManagerFn(router, db)
 
 	sortable, ok := manager.Handler().(rm.SortableHandler)
 	if ok {

server/testdb/data_stores_test.go

@@ -2,14 +2,14 @@ package testdb_test
 
 import (
 	"context"
+	"database/sql"
 	"testing"
 
 	"github.com/gorilla/mux"
 	"github.com/kubeshop/tracetest/server/model"
 	"github.com/kubeshop/tracetest/server/resourcemanager"
 	rmtests "github.com/kubeshop/tracetest/server/resourcemanager/testutil"
 	"github.com/kubeshop/tracetest/server/testdb"
-	"github.com/kubeshop/tracetest/server/testmock"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/collector/config/configgrpc"
@@ -177,16 +177,15 @@ func TestDataStoreProvisioner(t *testing.T) {
 	rmtests.TestResourceType(t, rmtests.ResourceTypeTest{
 		ResourceTypeSingular: testdb.DataStoreResourceName,
 		ResourceTypePlural:   testdb.DataStoreResourceName,
-		RegisterManagerFn: func(router *mux.Router) resourcemanager.Manager {
-			db := testmock.CreateMigratedDatabase()
+		RegisterManagerFn: func(router *mux.Router, db *sql.DB) resourcemanager.Manager {
 			dsRepo, err := testdb.Postgres(testdb.WithDB(db))
 			require.NoError(t, err)
 
 			manager := resourcemanager.New[testdb.DataStoreResource](
 				testdb.DataStoreResourceName,
 				testdb.DataStoreResourceNamePlural,
 				testdb.NewDataStoreResourceProvisioner(dsRepo),
-				// this resource exists only for provisiooning at the moment``
+				// this resource exists only for provisioning at the moment
 				resourcemanager.WithOperations(resourcemanager.OperationNoop),
 			)
 			manager.RegisterRoutes(router)