Improve tests cleanup

cmd/ferretdb/main.go

@@ -210,6 +210,7 @@ func run() {
 	logger := setupLogger(uuid)
 
 	ctx, stop := notifyAppTermination(context.Background())
+
 	go func() {
 		<-ctx.Done()
 		logger.Info("Stopping...")

cmd/ferretdb/main_unix.go

@@ -23,6 +23,7 @@ import (
 	"golang.org/x/sys/unix"
 )
 
+// notifyAppTermination installs a signal handler that cancels the context.
 func notifyAppTermination(parent context.Context) (context.Context, context.CancelFunc) {
 	return signal.NotifyContext(parent, unix.SIGTERM, unix.SIGINT)
 }

cmd/ferretdb/main_windows.go

@@ -22,6 +22,7 @@ import (
 	"golang.org/x/sys/windows"
 )
 
+// notifyAppTermination installs a signal handler that cancels the context.
 func notifyAppTermination(parent context.Context) (context.Context, context.CancelFunc) {
 	return signal.NotifyContext(parent, windows.SIGTERM, windows.SIGINT, os.Interrupt)
 }

internal/handlers/pg/pgdb/query_test.go

@@ -87,6 +87,7 @@ func TestQueryDocuments(t *testing.T) {
 
 			tx, err := pool.Begin(ctx)
 			require.NoError(t, err)
+			defer tx.Rollback(ctx)
 
 			for _, doc := range tc.documents {
 				require.NoError(t, InsertDocument(ctx, tx, databaseName, tc.collection, doc))
@@ -117,6 +118,7 @@ func TestQueryDocuments(t *testing.T) {
 	t.Run("cancel_context", func(t *testing.T) {
 		tx, err := pool.Begin(ctx)
 		require.NoError(t, err)
+		defer tx.Rollback(ctx)
 
 		for i := 1; i <= FetchedChannelBufSize*FetchedSliceCapacity+1; i++ {
 			require.NoError(t, InsertDocument(ctx, tx, databaseName, collectionName+"_cancel",
@@ -125,7 +127,7 @@ func TestQueryDocuments(t *testing.T) {
 		}
 
 		sp := &SQLParam{DB: databaseName, Collection: collectionName + "_cancel"}
-		ctx, cancel := context.WithCancel(context.Background())
+		ctx, cancel := context.WithCancel(ctx)
 		fetchedChan, err := pool.QueryDocuments(ctx, tx, sp)
 		cancel()
 		require.NoError(t, err)
@@ -149,9 +151,10 @@ func TestQueryDocuments(t *testing.T) {
 	t.Run("non-existing_collection", func(t *testing.T) {
 		tx, err := pool.Begin(ctx)
 		require.NoError(t, err)
+		defer tx.Rollback(ctx)
 
 		sp := &SQLParam{DB: databaseName, Collection: collectionName + "_non-existing"}
-		fetchedChan, err := pool.QueryDocuments(context.Background(), tx, sp)
+		fetchedChan, err := pool.QueryDocuments(ctx, tx, sp)
 		require.NoError(t, err)
 		res, ok := <-fetchedChan
 		require.False(t, ok)

internal/util/testutil/testutil.go

@@ -24,6 +24,13 @@ import (
 func Ctx(tb testing.TB) context.Context {
 	tb.Helper()
 
-	// TODO handle signals to stop tests gracefully
-	return context.Background()
+	ctx, stop := notifyTestsTermination(context.Background())
+
+	go func() {
+		<-ctx.Done()
+		tb.Log("Stopping...")
+		stop()
+	}()
+
+	return ctx
 }

internal/util/testutil/testutil_unix.go

@@ -0,0 +1,29 @@
+// Copyright 2021 FerretDB Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build !windows
+
+package testutil
+
+import (
+	"context"
+	"os/signal"
+
+	"golang.org/x/sys/unix"
+)
+
+// notifyTestsTermination installs a signal handler that cancels the context.
+func notifyTestsTermination(parent context.Context) (context.Context, context.CancelFunc) {
+	return signal.NotifyContext(parent, unix.SIGTERM, unix.SIGINT)
+}

internal/util/testutil/testutil_windows.go

@@ -0,0 +1,28 @@
+// Copyright 2021 FerretDB Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package testutil
+
+import (
+	"context"
+	"os"
+	"os/signal"
+
+	"golang.org/x/sys/windows"
+)
+
+// notifyTestsTermination installs a signal handler that cancels the context.
+func notifyTestsTermination(parent context.Context) (context.Context, context.CancelFunc) {
+	return signal.NotifyContext(parent, windows.SIGTERM, windows.SIGINT, os.Interrupt)
+}