fix flaky test with global session initializer

Author: agoncear-mwb

chdb/session_test.go

@@ -1,18 +1,47 @@
 package chdb
 
 import (
+	"fmt"
 	"os"
 	"path/filepath"
 	"testing"
 )
 
-// TestNewSession tests the creation of a new session.
-func TestNewSession(t *testing.T) {
-	session, err := NewSession()
+var (
+	session *Session
+)
+
+func globalSetup() error {
+	sess, err := NewSession()
 	if err != nil {
-		t.Fatalf("Failed to create new session: %s", err)
+		return err
 	}
-	defer session.Cleanup()
+	session = sess
+	return nil
+}
+
+func globalTeardown() {
+	session.Cleanup()
+	session.Close()
+}
+
+func TestMain(m *testing.M) {
+	if err := globalSetup(); err != nil {
+		fmt.Println("Global setup failed:", err)
+		os.Exit(1)
+	}
+	// Run all tests.
+	exitCode := m.Run()
+
+	// Global teardown: clean up any resources here.
+	globalTeardown()
+
+	// Exit with the code returned by m.Run().
+	os.Exit(exitCode)
+}
+
+// TestNewSession tests the creation of a new session.
+func TestNewSession(t *testing.T) {
 
 	// Check if the session directory exists
 	if _, err := os.Stat(session.Path()); os.IsNotExist(err) {
@@ -25,59 +54,30 @@ func TestNewSession(t *testing.T) {
 	}
 }
 
-// TestSessionClose tests the Close method of the session.
-func TestSessionClose(t *testing.T) {
-	session, _ := NewSession()
-	defer session.Cleanup() // Cleanup in case Close fails
+// This test is currently flaky because of this: https://github.com/chdb-io/chdb/pull/299/commits/91b0aedd8c17e74a4bb213e885d89cc9a77c99ad
+func TestQuery(t *testing.T) {
 
-	// Close the session
-	session.Close()
+	// time.Sleep(time.Second * 5)
 
-	// Check if the session directory has been removed
-	if _, err := os.Stat(session.Path()); !os.IsNotExist(err) {
-		t.Errorf("Session directory should be removed after Close: %s", session.Path())
+	_, err := session.Query("CREATE TABLE IF NOT EXISTS TestQuery (id UInt32) ENGINE = MergeTree() ORDER BY id;")
+	if err != nil {
+		t.Fatal(err)
 	}
-}
 
-// TestSessionCleanup tests the Cleanup method of the session.
-func TestSessionCleanup(t *testing.T) {
-	session, _ := NewSession()
-
-	// Cleanup the session
-	session.Cleanup()
+	_, err = session.Query("INSERT INTO TestQuery VALUES (1), (2), (3);")
+	if err != nil {
+		t.Fatal(err)
+	}
+	ret, err := session.Query("SELECT * FROM TestQuery;")
+	if err != nil {
+		t.Fatalf("Query failed: %s", err)
+	}
 
-	// Check if the session directory has been removed
-	if _, err := os.Stat(session.Path()); !os.IsNotExist(err) {
-		t.Errorf("Session directory should be removed after Cleanup: %s", session.Path())
+	if string(ret.Buf()) != "1\n2\n3\n" {
+		t.Fatalf("Query result should be 1\n2\n3\n, got %s", string(ret.Buf()))
 	}
 }
 
-// This test is currently flaky because of this: https://github.com/chdb-io/chdb/pull/299/commits/91b0aedd8c17e74a4bb213e885d89cc9a77c99ad
-// func TestQuery(t *testing.T) {
-
-// 	session, _ := NewSession()
-// 	defer session.Cleanup()
-// 	// time.Sleep(time.Second * 5)
-
-// 	_, err := session.Query("CREATE TABLE IF NOT EXISTS TestQuery (id UInt32) ENGINE = MergeTree() ORDER BY id;")
-// 	if err != nil {
-// 		t.Fatal(err)
-// 	}
-
-// 	_, err = session.Query("INSERT INTO TestQuery VALUES (1), (2), (3);")
-// 	if err != nil {
-// 		t.Fatal(err)
-// 	}
-// 	ret, err := session.Query("SELECT * FROM TestQuery;")
-// 	if err != nil {
-// 		t.Fatalf("Query failed: %s", err)
-// 	}
-
-// 	if string(ret.Buf()) != "1\n2\n3\n" {
-// 		t.Fatalf("Query result should be 1\n2\n3\n, got %s", string(ret.Buf()))
-// 	}
-// }
-
 func TestSessionPathAndIsTemp(t *testing.T) {
 	// Create a new session and check its Path and IsTemp
 	session, _ := NewSession()
@@ -105,3 +105,30 @@ func TestSessionPathAndIsTemp(t *testing.T) {
 		t.Errorf("Session should not be temporary")
 	}
 }
+
+// TestSessionClose tests the Close method of the session.
+func TestSessionClose(t *testing.T) {
+	session, _ := NewSession()
+	defer session.Cleanup() // Cleanup in case Close fails
+
+	// Close the session
+	session.Close()
+
+	// Check if the session directory has been removed
+	if _, err := os.Stat(session.Path()); !os.IsNotExist(err) {
+		t.Errorf("Session directory should be removed after Close: %s", session.Path())
+	}
+}
+
+// TestSessionCleanup tests the Cleanup method of the session.
+func TestSessionCleanup(t *testing.T) {
+	session, _ := NewSession()
+
+	// Cleanup the session
+	session.Cleanup()
+
+	// Check if the session directory has been removed
+	if _, err := os.Stat(session.Path()); !os.IsNotExist(err) {
+		t.Errorf("Session directory should be removed after Cleanup: %s", session.Path())
+	}
+}