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/handlers/pg/pgdb/query_test.go b/internal/handlers/pg/pgdb/query_test.go index 3739cf6f9b23..5528eea8b4e1 100644 --- a/internal/handlers/pg/pgdb/query_test.go +++ b/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) 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) +}