chore: use t.Cleanup() for integration tests cleanup (#382)

Use `t.Cleanup` to make sure the cleanup logic is registered right after the resource is created and always gets executed when the test finishes.
DataDog · Nov 7, 2024 · a4ba006 · a4ba006
1 parent 6b2ab98
commit a4ba006
Show file tree

Hide file tree

Showing 56 changed files with 128 additions and 300 deletions.
diff --git a/_integration-tests/tests/99designs.gqlgen/gqlgen.go b/_integration-tests/tests/99designs.gqlgen/gqlgen.go
@@ -58,8 +58,6 @@ func (tc *TestCase) Run(t *testing.T) {
 	}, resp)
 }
 
-func (*TestCase) Teardown(*testing.T) {}
-
 func (*TestCase) ExpectedTraces() trace.Traces {
 	return trace.Traces{
 		{

diff --git a/_integration-tests/tests/aws.v1/aws.go b/_integration-tests/tests/aws.v1/aws.go
@@ -8,10 +8,8 @@
 package awsv1
 
 import (
-	"context"
 	"fmt"
 	"testing"
-	"time"
 
 	"datadoghq.dev/orchestrion/_integration-tests/utils"
 	"datadoghq.dev/orchestrion/_integration-tests/validator/trace"
@@ -47,12 +45,6 @@ func (tc *TestCase) Run(t *testing.T) {
 	require.NoError(t, err)
 }
 
-func (tc *TestCase) Teardown(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancel()
-	require.NoError(t, tc.server.Terminate(ctx))
-}
-
 func (*TestCase) ExpectedTraces() trace.Traces {
 	return trace.Traces{
 		{

diff --git a/_integration-tests/tests/aws.v2/base.go b/_integration-tests/tests/aws.v2/base.go
@@ -34,12 +34,6 @@ func (b *base) setup(t *testing.T) {
 	b.server, b.host, b.port = utils.StartDynamoDBTestContainer(t)
 }
 
-func (b *base) teardown(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancel()
-	require.NoError(t, b.server.Terminate(ctx))
-}
-
 func (b *base) run(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
 	defer cancel()

diff --git a/_integration-tests/tests/aws.v2/load_default_config.go b/_integration-tests/tests/aws.v2/load_default_config.go
@@ -39,10 +39,6 @@ func (tc *TestCaseLoadDefaultConfig) Run(t *testing.T) {
 	tc.base.run(t)
 }
 
-func (tc *TestCaseLoadDefaultConfig) Teardown(t *testing.T) {
-	tc.base.teardown(t)
-}
-
 func (tc *TestCaseLoadDefaultConfig) ExpectedTraces() trace.Traces {
 	return tc.base.expectedTraces()
 }
diff --git a/_integration-tests/tests/aws.v2/new_config.go b/_integration-tests/tests/aws.v2/new_config.go
@@ -34,10 +34,6 @@ func (tc *TestCaseNewConfig) Run(t *testing.T) {
 	tc.base.run(t)
 }
 
-func (tc *TestCaseNewConfig) Teardown(t *testing.T) {
-	tc.base.teardown(t)
-}
-
 func (tc *TestCaseNewConfig) ExpectedTraces() trace.Traces {
 	return tc.base.expectedTraces()
 }
diff --git a/_integration-tests/tests/aws.v2/struct_literal.go b/_integration-tests/tests/aws.v2/struct_literal.go
@@ -34,10 +34,6 @@ func (tc *TestCaseStructLiteral) Run(t *testing.T) {
 	tc.base.run(t)
 }
 
-func (tc *TestCaseStructLiteral) Teardown(t *testing.T) {
-	tc.base.teardown(t)
-}
-
 func (tc *TestCaseStructLiteral) ExpectedTraces() trace.Traces {
 	return tc.base.expectedTraces()
 }
diff --git a/_integration-tests/tests/aws.v2/struct_literal_ptr.go b/_integration-tests/tests/aws.v2/struct_literal_ptr.go
@@ -35,10 +35,6 @@ func (tc *TestCaseStructLiteralPtr) Run(t *testing.T) {
 	tc.base.run(t)
 }
 
-func (tc *TestCaseStructLiteralPtr) Teardown(t *testing.T) {
-	tc.base.teardown(t)
-}
-
 func (tc *TestCaseStructLiteralPtr) ExpectedTraces() trace.Traces {
 	return tc.base.expectedTraces()
 }
diff --git a/_integration-tests/tests/chi.v5/chi.go b/_integration-tests/tests/chi.v5/chi.go
@@ -39,6 +39,11 @@ func (tc *TestCase) Setup(t *testing.T) {
 	})
 
 	go func() { assert.ErrorIs(t, tc.Server.ListenAndServe(), http.ErrServerClosed) }()
+	t.Cleanup(func() {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+		assert.NoError(t, tc.Server.Shutdown(ctx))
+	})
 }
 
 func (tc *TestCase) Run(t *testing.T) {
@@ -47,13 +52,6 @@ func (tc *TestCase) Run(t *testing.T) {
 	require.Equal(t, http.StatusOK, resp.StatusCode)
 }
 
-func (tc *TestCase) Teardown(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
-	defer cancel()
-
-	require.NoError(t, tc.Server.Shutdown(ctx))
-}
-
 func (tc *TestCase) ExpectedTraces() trace.Traces {
 	return trace.Traces{
 		{

diff --git a/_integration-tests/tests/confluent-kafka-go.v1/kafka.go b/_integration-tests/tests/confluent-kafka-go.v1/kafka.go
@@ -8,7 +8,6 @@
 package kafka
 
 import (
-	"context"
 	"strings"
 	"testing"
 	"time"
@@ -104,13 +103,6 @@ func (tc *TestCase) consumeMessage(t *testing.T) {
 	require.Equal(t, "key2", string(m.Key))
 }
 
-func (tc *TestCase) Teardown(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
-	defer cancel()
-
-	require.NoError(t, tc.container.Terminate(ctx))
-}
-
 func (*TestCase) ExpectedTraces() trace.Traces {
 	return trace.Traces{
 		{

diff --git a/_integration-tests/tests/confluent-kafka-go.v1/skip_windows.go b/_integration-tests/tests/confluent-kafka-go.v1/skip_windows.go
@@ -20,7 +20,6 @@ func (skip) Setup(t *testing.T) {
 }
 
 func (skip) Run(t *testing.T)             {}
-func (skip) Teardown(t *testing.T)        {}
 func (skip) ExpectedTraces() trace.Traces { return nil }
 
 type TestCase = skip
diff --git a/_integration-tests/tests/confluent-kafka-go.v2/kafka.go b/_integration-tests/tests/confluent-kafka-go.v2/kafka.go
@@ -8,7 +8,6 @@
 package kafka
 
 import (
-	"context"
 	"strings"
 	"testing"
 	"time"
@@ -104,13 +103,6 @@ func (tc *TestCase) consumeMessage(t *testing.T) {
 	require.Equal(t, "key2", string(m.Key))
 }
 
-func (tc *TestCase) Teardown(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
-	defer cancel()
-
-	require.NoError(t, tc.container.Terminate(ctx))
-}
-
 func (*TestCase) ExpectedTraces() trace.Traces {
 	return trace.Traces{
 		{

diff --git a/_integration-tests/tests/confluent-kafka-go.v2/skip_windows.go b/_integration-tests/tests/confluent-kafka-go.v2/skip_windows.go
@@ -20,7 +20,6 @@ func (skip) Setup(t *testing.T) {
 }
 
 func (skip) Run(t *testing.T)             {}
-func (skip) Teardown(t *testing.T)        {}
 func (skip) ExpectedTraces() trace.Traces { return nil }
 
 type TestCase = skip
diff --git a/_integration-tests/tests/dd-span/ddspan.go b/_integration-tests/tests/dd-span/ddspan.go
@@ -30,8 +30,6 @@ func (*TestCase) Run(t *testing.T) {
 	_, _ = spanFromHTTPRequest(req)
 }
 
-func (*TestCase) Teardown(*testing.T) {}
-
 //dd:span foo:bar
 func spanFromHTTPRequest(*http.Request) (string, error) {
 	return tagSpecificSpan()

diff --git a/_integration-tests/tests/echo.v4/echo.go b/_integration-tests/tests/echo.v4/echo.go
@@ -36,6 +36,11 @@ func (tc *TestCase) Setup(t *testing.T) {
 	tc.addr = "127.0.0.1:" + utils.GetFreePort(t)
 
 	go func() { assert.ErrorIs(t, tc.Echo.Start(tc.addr), http.ErrServerClosed) }()
+	t.Cleanup(func() {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+		require.NoError(t, tc.Echo.Shutdown(ctx))
+	})
 }
 
 func (tc *TestCase) Run(t *testing.T) {
@@ -44,13 +49,6 @@ func (tc *TestCase) Run(t *testing.T) {
 	require.Equal(t, http.StatusOK, resp.StatusCode)
 }
 
-func (tc *TestCase) Teardown(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
-	defer cancel()
-
-	require.NoError(t, tc.Echo.Shutdown(ctx))
-}
-
 func (tc *TestCase) ExpectedTraces() trace.Traces {
 	return trace.Traces{
 		{

diff --git a/_integration-tests/tests/fiber.v2/fiber.go b/_integration-tests/tests/fiber.v2/fiber.go
@@ -30,6 +30,9 @@ func (tc *TestCase) Setup(t *testing.T) {
 	tc.addr = "127.0.0.1:" + utils.GetFreePort(t)
 
 	go func() { assert.NoError(t, tc.App.Listen(tc.addr)) }()
+	t.Cleanup(func() {
+		assert.NoError(t, tc.App.ShutdownWithTimeout(time.Second))
+	})
 }
 
 func (tc *TestCase) Run(t *testing.T) {
@@ -38,10 +41,6 @@ func (tc *TestCase) Run(t *testing.T) {
 	require.Equal(t, http.StatusOK, resp.StatusCode)
 }
 
-func (tc *TestCase) Teardown(t *testing.T) {
-	require.NoError(t, tc.App.ShutdownWithTimeout(time.Second))
-}
-
 func (tc *TestCase) ExpectedTraces() trace.Traces {
 	return trace.Traces{
 		{

diff --git a/_integration-tests/tests/gcp_pubsub/gcp_pubsub.go b/_integration-tests/tests/gcp_pubsub/gcp_pubsub.go
@@ -51,6 +51,7 @@ func (tc *TestCase) Setup(t *testing.T) {
 		utils.WithTestLogConsumer(t),
 	)
 	utils.AssertTestContainersError(t, err)
+	utils.RegisterContainerCleanup(t, tc.container)
 
 	projectID := tc.container.Settings.ProjectID
 
@@ -60,6 +61,9 @@ func (tc *TestCase) Setup(t *testing.T) {
 
 	tc.client, err = pubsub.NewClient(ctx, projectID, option.WithGRPCConn(conn))
 	require.NoError(t, err)
+	t.Cleanup(func() {
+		assert.NoError(t, tc.client.Close())
+	})
 
 	topic, err := tc.client.CreateTopic(ctx, testTopic)
 	require.NoError(t, err)
@@ -111,14 +115,6 @@ func (tc *TestCase) Run(t *testing.T) {
 	tc.receiveMessage(t)
 }
 
-func (tc *TestCase) Teardown(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
-	defer cancel()
-
-	require.NoError(t, tc.client.Close())
-	require.NoError(t, tc.container.Terminate(ctx))
-}
-
 func (tc *TestCase) ExpectedTraces() trace.Traces {
 	return trace.Traces{
 		{

diff --git a/_integration-tests/tests/gin/gin.go b/_integration-tests/tests/gin/gin.go
@@ -36,6 +36,11 @@ func (tc *TestCase) Setup(t *testing.T) {
 	engine.GET("/ping", func(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"message": "pong"}) })
 
 	go func() { assert.ErrorIs(t, tc.Server.ListenAndServe(), http.ErrServerClosed) }()
+	t.Cleanup(func() {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+		assert.NoError(t, tc.Server.Shutdown(ctx))
+	})
 }
 
 func (tc *TestCase) Run(t *testing.T) {
@@ -44,13 +49,6 @@ func (tc *TestCase) Run(t *testing.T) {
 	require.Equal(t, http.StatusOK, resp.StatusCode)
 }
 
-func (tc *TestCase) Teardown(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
-	defer cancel()
-
-	require.NoError(t, tc.Server.Shutdown(ctx))
-}
-
 func (tc *TestCase) ExpectedTraces() trace.Traces {
 	return trace.Traces{
 		{

diff --git a/_integration-tests/tests/go-elasticsearch/base.go b/_integration-tests/tests/go-elasticsearch/base.go
@@ -18,7 +18,6 @@ import (
 
 	"datadoghq.dev/orchestrion/_integration-tests/utils"
 	"datadoghq.dev/orchestrion/_integration-tests/validator/trace"
-	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/testcontainers/testcontainers-go"
 	testelasticsearch "github.com/testcontainers/testcontainers-go/modules/elasticsearch"
@@ -48,6 +47,7 @@ func (b *base) Setup(t *testing.T, image string, newClient func(addr string, caC
 		testcontainers.WithWaitStrategyAndDeadline(time.Minute, wait.ForLog(`.*("message":\s?"started(\s|")?.*|]\sstarted\n)`).AsRegexp()),
 	)
 	utils.AssertTestContainersError(t, err)
+	utils.RegisterContainerCleanup(t, b.container)
 
 	b.client, err = newClient(b.container.Settings.Address, b.container.Settings.CACert)
 	require.NoError(t, err)
@@ -66,13 +66,6 @@ func (b *base) Run(t *testing.T, doRequest func(t *testing.T, client esClient, b
 	doRequest(t, b.client, bytes.NewReader(data))
 }
 
-func (b *base) Teardown(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
-	defer cancel()
-
-	assert.NoError(t, b.container.Terminate(ctx))
-}
-
 func (*base) ExpectedTraces() trace.Traces {
 	return trace.Traces{
 		{

diff --git a/_integration-tests/tests/go-elasticsearch/skip_windows.go b/_integration-tests/tests/go-elasticsearch/skip_windows.go
@@ -20,7 +20,6 @@ func (skip) Setup(t *testing.T) {
 }
 
 func (skip) Run(t *testing.T)             {}
-func (skip) Teardown(t *testing.T)        {}
 func (skip) ExpectedTraces() trace.Traces { return nil }
 
 type TestCaseV6 = skip

diff --git a/_integration-tests/tests/go-redis.v0/go-redis.go b/_integration-tests/tests/go-redis.v0/go-redis.go
@@ -11,7 +11,6 @@ import (
 	"context"
 	"fmt"
 	"testing"
-	"time"
 
 	"datadoghq.dev/orchestrion/_integration-tests/utils"
 	"datadoghq.dev/orchestrion/_integration-tests/validator/trace"
@@ -40,6 +39,9 @@ func (tc *TestCase) Setup(t *testing.T) {
 	tc.server = container
 
 	tc.Client = redis.NewClient(&redis.Options{Addr: addr})
+	t.Cleanup(func() {
+		assert.NoError(t, tc.Client.Close())
+	})
 }
 
 func (tc *TestCase) Run(t *testing.T) {
@@ -50,16 +52,6 @@ func (tc *TestCase) Run(t *testing.T) {
 	require.NoError(t, tc.Client.WithContext(ctx).Get(tc.key).Err())
 }
 
-func (tc *TestCase) Teardown(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
-	defer cancel()
-
-	assert.NoError(t, tc.Client.Close())
-	if tc.server != nil && assert.NoError(t, tc.server.Terminate(ctx)) {
-		tc.server = nil
-	}
-}
-
 func (tc *TestCase) ExpectedTraces() trace.Traces {
 	return trace.Traces{
 		{

diff --git a/_integration-tests/tests/go-redis.v7/go-redis.go b/_integration-tests/tests/go-redis.v7/go-redis.go
@@ -11,7 +11,6 @@ import (
 	"context"
 	"fmt"
 	"testing"
-	"time"
 
 	"datadoghq.dev/orchestrion/_integration-tests/utils"
 	"datadoghq.dev/orchestrion/_integration-tests/validator/trace"
@@ -40,6 +39,9 @@ func (tc *TestCase) Setup(t *testing.T) {
 	tc.server = container
 
 	tc.Client = redis.NewClient(&redis.Options{Addr: addr})
+	t.Cleanup(func() {
+		assert.NoError(t, tc.Client.Close())
+	})
 }
 
 func (tc *TestCase) Run(t *testing.T) {
@@ -50,16 +52,6 @@ func (tc *TestCase) Run(t *testing.T) {
 	require.NoError(t, tc.Client.WithContext(ctx).Get(tc.key).Err())
 }
 
-func (tc *TestCase) Teardown(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
-	defer cancel()
-
-	assert.NoError(t, tc.Client.Close())
-	if tc.server != nil && assert.NoError(t, tc.server.Terminate(ctx)) {
-		tc.server = nil
-	}
-}
-
 func (tc *TestCase) ExpectedTraces() trace.Traces {
 	return trace.Traces{
 		{