Ensure test container termination (#1274)

* Ensure test container termination * make #1229 more resilient * make #1229 more resilent
ClickHouse · Apr 15, 2024 · 6fdbf7d · 6fdbf7d
1 parent aa43b3f
commit 6fdbf7d
Show file tree

Hide file tree

Showing 6 changed files with 57 additions and 107 deletions.
diff --git a/tests/issues/1127_test.go b/tests/issues/1127_test.go
@@ -13,12 +13,7 @@ import (
 
 func Test1127(t *testing.T) {
 	var (
-		conn, err = clickhouse_tests.GetConnection("issues", clickhouse.Settings{
-			"max_execution_time":             60,
-			"allow_experimental_object_type": true,
-		}, nil, &clickhouse.Compression{
-			Method: clickhouse.CompressionLZ4,
-		})
+		conn, err = clickhouse_tests.GetConnection("issues", nil, nil, nil)
 	)
 	require.NoError(t, err)
 
@@ -30,11 +25,14 @@ func Test1127(t *testing.T) {
 		fmt.Println("log info: ", log)
 	}))
 
-	rows, err := conn.Query(ctx, "select throwIf(number = 1e6) from system.numbers settings max_block_size = 100")
+	rows, err := conn.Query(ctx, "select number, throwIf(number = 1e6) from system.numbers settings max_block_size = 100")
 	require.NoError(t, err)
-
 	defer rows.Close()
+
+	var number uint64
+	var throwIf uint8
 	for rows.Next() {
+		require.NoError(t, rows.Scan(&number, &throwIf))
 	}
 
 	assert.Error(t, rows.Err())

diff --git a/tests/issues/1229_test.go b/tests/issues/1229_test.go
@@ -48,7 +48,9 @@ func Test1229(t *testing.T) {
 			defer wg.Done()
 			withTimeoutCtx, cancel := context.WithTimeout(ctx, queryTimeout)
 			defer cancel()
-			_, _ = conn.Query(withTimeoutCtx, selectQuery)
+			rows, err := conn.Query(withTimeoutCtx, selectQuery)
+			require.NoError(t, err)
+			require.NoError(t, rows.Close())
 		}()
 	}
 

diff --git a/tests/issues/main_test.go b/tests/issues/main_test.go
@@ -18,45 +18,16 @@
 package issues
 
 import (
-	"context"
-	"fmt"
-	clickhouse_tests "github.com/ClickHouse/clickhouse-go/v2/tests"
-	"math/rand"
 	"os"
-	"strconv"
 	"testing"
-	"time"
+
+	clickhouse_tests "github.com/ClickHouse/clickhouse-go/v2/tests"
 )
 
 const testSet string = "issues"
 
 func TestMain(m *testing.M) {
-	seed := time.Now().UnixNano()
-	fmt.Printf("using random seed %d for %s\n", seed, testSet)
-	rand.Seed(seed)
-	useDocker, err := strconv.ParseBool(clickhouse_tests.GetEnv("CLICKHOUSE_USE_DOCKER", "true"))
-	if err != nil {
-		panic(err)
-	}
-	var env clickhouse_tests.ClickHouseTestEnvironment
-	switch useDocker {
-	case true:
-		env, err = clickhouse_tests.CreateClickHouseTestEnvironment(testSet)
-		if err != nil {
-			panic(err)
-		}
-		defer env.Container.Terminate(context.Background()) //nolint
-	case false:
-		env, err = clickhouse_tests.GetExternalTestEnvironment(testSet)
-		if err != nil {
-			panic(err)
-		}
-	}
-	clickhouse_tests.SetTestEnvironment(testSet, env)
-	if err := clickhouse_tests.CreateDatabase(testSet); err != nil {
-		panic(err)
-	}
-	os.Exit(m.Run())
+	os.Exit(clickhouse_tests.Runtime(m, testSet))
 }
 
 func GetIssuesTestEnvironment() (clickhouse_tests.ClickHouseTestEnvironment, error) {

diff --git a/tests/main_test.go b/tests/main_test.go
@@ -18,47 +18,18 @@
 package tests
 
 import (
-	"context"
 	"crypto/tls"
-	"fmt"
-	"github.com/ClickHouse/clickhouse-go/v2"
-	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
-	"math/rand"
 	"os"
-	"strconv"
 	"testing"
-	"time"
+
+	"github.com/ClickHouse/clickhouse-go/v2"
+	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
 )
 
 const testSet string = "native"
 
 func TestMain(m *testing.M) {
-	seed := time.Now().UnixNano()
-	fmt.Printf("using random seed %d for %s tests\n", seed, testSet)
-	rand.Seed(seed)
-	useDocker, err := strconv.ParseBool(GetEnv("CLICKHOUSE_USE_DOCKER", "true"))
-	if err != nil {
-		panic(err)
-	}
-	var env ClickHouseTestEnvironment
-	switch useDocker {
-	case true:
-		env, err = CreateClickHouseTestEnvironment(testSet)
-		if err != nil {
-			panic(err)
-		}
-		defer env.Container.Terminate(context.Background()) //nolint
-	case false:
-		env, err = GetExternalTestEnvironment(testSet)
-		if err != nil {
-			panic(err)
-		}
-	}
-	SetTestEnvironment(testSet, env)
-	if err := CreateDatabase(testSet); err != nil {
-		panic(err)
-	}
-	os.Exit(m.Run())
+	os.Exit(Runtime(m, testSet))
 }
 
 func GetNativeTestEnvironment() (ClickHouseTestEnvironment, error) {

diff --git a/tests/std/main_test.go b/tests/std/main_test.go
@@ -18,49 +18,20 @@
 package std
 
 import (
-	"context"
 	"crypto/tls"
 	"database/sql"
-	"fmt"
-	"github.com/ClickHouse/clickhouse-go/v2"
-	clickhouse_tests "github.com/ClickHouse/clickhouse-go/v2/tests"
-	"math/rand"
 	"net/url"
 	"os"
-	"strconv"
 	"testing"
-	"time"
+
+	"github.com/ClickHouse/clickhouse-go/v2"
+	clickhouse_tests "github.com/ClickHouse/clickhouse-go/v2/tests"
 )
 
 const testSet string = "std"
 
 func TestMain(m *testing.M) {
-	seed := time.Now().UnixNano()
-	fmt.Printf("using random seed %d for %s tests\n", seed, testSet)
-	rand.Seed(seed)
-	useDocker, err := strconv.ParseBool(clickhouse_tests.GetEnv("CLICKHOUSE_USE_DOCKER", "true"))
-	if err != nil {
-		panic(err)
-	}
-	var env clickhouse_tests.ClickHouseTestEnvironment
-	switch useDocker {
-	case true:
-		env, err = clickhouse_tests.CreateClickHouseTestEnvironment(testSet)
-		if err != nil {
-			panic(err)
-		}
-		defer env.Container.Terminate(context.Background()) //nolint
-	case false:
-		env, err = clickhouse_tests.GetExternalTestEnvironment(testSet)
-		if err != nil {
-			panic(err)
-		}
-	}
-	clickhouse_tests.SetTestEnvironment(testSet, env)
-	if err := clickhouse_tests.CreateDatabase(testSet); err != nil {
-		panic(err)
-	}
-	os.Exit(m.Run())
+	os.Exit(clickhouse_tests.Runtime(m, testSet))
 }
 
 func GetStdDSNConnection(protocol clickhouse.Protocol, secure bool, opts url.Values) (*sql.DB, error) {

diff --git a/tests/utils.go b/tests/utils.go
@@ -759,3 +759,40 @@ func optionsToDSN(o *clickhouse.Options) string {
 
 	return u.String()
 }
+
+func Runtime(m *testing.M, ts string) (exitCode int) {
+	seed := time.Now().UnixNano()
+	rand.Seed(seed)
+	fmt.Printf("using random seed %d for %s tests\n", seed, ts)
+
+	useDocker, err := strconv.ParseBool(GetEnv("CLICKHOUSE_USE_DOCKER", "true"))
+	if err != nil {
+		panic(err)
+	}
+
+	var env ClickHouseTestEnvironment
+	switch useDocker {
+	case true:
+		env, err = CreateClickHouseTestEnvironment(ts)
+		if err != nil {
+			panic(err)
+		}
+		defer func() {
+			if err := env.Container.Terminate(context.Background()); err != nil {
+				panic(err)
+			}
+		}() //nolint
+	case false:
+		env, err = GetExternalTestEnvironment(ts)
+		if err != nil {
+			panic(err)
+		}
+	}
+
+	SetTestEnvironment(ts, env)
+	if err := CreateDatabase(ts); err != nil {
+		panic(err)
+	}
+
+	return m.Run()
+}