From 7ce08bc788ede6cff55736fd83ef119c08b449fd Mon Sep 17 00:00:00 2001
From: Alexander Todorov <atodorov@redhat.com>
Date: Fri, 24 Apr 2020 12:48:39 +0200
Subject: [PATCH] tests: Make sure TestMain cleans up after itself

os.Exit() doesn't execute defer-ed functions, OTOH we call panic()
in case of setup errors, which does run defer-ed functions. So move
the actual test setup and execution in a separate function which will
execute all defer-ed functions and return the exit code from the
test suite to TestMain.
---
 internal/client/integration_test.go | 28 +++++++++++++++++-----------
 internal/client/unit_test.go        | 10 +++++++---
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/internal/client/integration_test.go b/internal/client/integration_test.go
index a1468544d..20bd62b44 100644
--- a/internal/client/integration_test.go
+++ b/internal/client/integration_test.go
@@ -19,29 +19,35 @@ var testState *TestState
 
 // Setup the socket to use for running the tests
 // Also makes sure there is a running server to test against
-func TestMain(m *testing.M) {
+func executeTests(m *testing.M) int {
 	var err error
 	testState, err = setUpTestState("/run/weldr/api.socket", 60*time.Second, false)
 	if err != nil {
 		fmt.Printf("ERROR: Test setup failed: %s\n", err)
-		os.Exit(1)
+		panic(err)
 	}
 
 	// Setup the test repo
 	dir, err := test.SetUpTemporaryRepository()
 	if err != nil {
 		fmt.Printf("ERROR: Test repo setup failed: %s\n", err)
-		os.Exit(1)
+		panic(err)
 	}
+
+	// Cleanup after the tests
+	defer func() {
+		err := test.TearDownTemporaryRepository(dir)
+		if err != nil {
+			fmt.Printf("ERROR: Failed to clean up temporary repository: %s\n", err)
+		}
+	}()
+
 	testState.repoDir = dir
 
 	// Run the tests
-	rc := m.Run()
-
-	// Cleanup after the tests
-	err = test.TearDownTemporaryRepository(dir)
-	if err != nil {
-		fmt.Printf("ERROR: Failed to clean up temporary repository: %s\n", err)
-	}
-	os.Exit(rc)
+	return m.Run()
+}
+
+func TestMain(m *testing.M) {
+	os.Exit(executeTests(m))
 }
diff --git a/internal/client/unit_test.go b/internal/client/unit_test.go
index d1cde6bf6..d2230cc50 100644
--- a/internal/client/unit_test.go
+++ b/internal/client/unit_test.go
@@ -23,7 +23,7 @@ import (
 // Hold test state to share between tests
 var testState *TestState
 
-func TestMain(m *testing.M) {
+func executeTests(m *testing.M) int {
 	// Setup the mocked server running on a temporary domain socket
 	tmpdir, err := ioutil.TempDir("", "client_test-")
 	if err != nil {
@@ -53,7 +53,7 @@ func TestMain(m *testing.M) {
 
 	go func() {
 		err := server.Serve(ln)
-		if err != nil {
+		if err != nil && err != http.ErrServerClosed {
 			panic(err)
 		}
 	}()
@@ -64,5 +64,9 @@ func TestMain(m *testing.M) {
 	}
 
 	// Run the tests
-	os.Exit(m.Run())
+	return m.Run()
+}
+
+func TestMain(m *testing.M) {
+	os.Exit(executeTests(m))
 }