From df6bb13741810df8a12b0039b1b9afb581985ece Mon Sep 17 00:00:00 2001
From: Alexey Palazhchenko <alexey.palazhchenko@ferretdb.io>
Date: Wed, 19 Oct 2022 17:58:15 +0400
Subject: [PATCH] Stop tests gracefully

---
 cmd/ferretdb/main.go                       |  1 +
 cmd/ferretdb/main_unix.go                  |  1 +
 cmd/ferretdb/main_windows.go               |  1 +
 internal/util/testutil/testutil.go         | 11 ++++++--
 internal/util/testutil/testutil_unix.go    | 29 ++++++++++++++++++++++
 internal/util/testutil/testutil_windows.go | 28 +++++++++++++++++++++
 6 files changed, 69 insertions(+), 2 deletions(-)
 create mode 100644 internal/util/testutil/testutil_unix.go
 create mode 100644 internal/util/testutil/testutil_windows.go

diff --git a/cmd/ferretdb/main.go b/cmd/ferretdb/main.go
index 89ffbd46f1ef..f23f3552d820 100644
--- a/cmd/ferretdb/main.go
+++ b/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...")
diff --git a/cmd/ferretdb/main_unix.go b/cmd/ferretdb/main_unix.go
index 0bd7a8731e43..b646d19fed56 100644
--- a/cmd/ferretdb/main_unix.go
+++ b/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)
 }
diff --git a/cmd/ferretdb/main_windows.go b/cmd/ferretdb/main_windows.go
index e513a3e558ab..a382597e59f3 100644
--- a/cmd/ferretdb/main_windows.go
+++ b/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)
 }
diff --git a/internal/util/testutil/testutil.go b/internal/util/testutil/testutil.go
index 024e870d6fae..a7e5a3ff7bf6 100644
--- a/internal/util/testutil/testutil.go
+++ b/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
 }
diff --git a/internal/util/testutil/testutil_unix.go b/internal/util/testutil/testutil_unix.go
new file mode 100644
index 000000000000..79b434900cdf
--- /dev/null
+++ b/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)
+}
diff --git a/internal/util/testutil/testutil_windows.go b/internal/util/testutil/testutil_windows.go
new file mode 100644
index 000000000000..97b2650e7073
--- /dev/null
+++ b/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)
+}