Fix flacky tests

internal/topic/topicwriterinternal/writer_reconnector_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopiccommon"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopicwriter"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawydb"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xatomic"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
@@ -360,14 +361,13 @@ func TestWriterImpl_WaitInit(t *testing.T) {
 		ctx, cancel := context.WithCancel(context.Background())
 
 		go func() {
+			// wait until w.WaitInit starts
+			time.Sleep(time.Millisecond)
 			cancel()
-			w.onWriterChange(&SingleStreamWriter{})
 		}()
 
 		_, err := w.WaitInit(ctx)
 		require.ErrorIs(t, err, ctx.Err())
-
-		require.True(t, isClosed(w.firstInitResponseProcessedChan))
 	})
 
 	t.Run("contextDeadlineErrorBeforeStart", func(t *testing.T) {
@@ -426,50 +426,53 @@ func TestWriterImpl_Reconnect(t *testing.T) {
 			connectionError error
 		}
 
-		newStream := func(name string, onSendInitCallback func()) *MockRawTopicWriterStream {
+		newStream := func(name string) *MockRawTopicWriterStream {
 			strm := NewMockRawTopicWriterStream(mc)
 			initReq := testCreateInitRequest(w)
 
 			streamClosed := make(empty.Chan)
 			strm.EXPECT().CloseSend().Do(func() {
+				t.Logf("closed stream: %v", name)
 				close(streamClosed)
 			})
 
 			strm.EXPECT().Send(&initReq).Do(func(_ interface{}) {
-				if onSendInitCallback != nil {
-					onSendInitCallback()
-				}
+				t.Logf("sent init request stream: %v", name)
 			})
 
-			strm.EXPECT().Recv().Return(&rawtopicwriter.InitResult{
+			strm.EXPECT().Recv().Do(func() {
+				t.Logf("receive init response stream: %v", name)
+			}).Return(&rawtopicwriter.InitResult{
 				ServerMessageMetadata: rawtopiccommon.ServerMessageMetadata{Status: rawydb.StatusSuccess},
 				SessionID:             name,
 			}, nil)
 
 			strm.EXPECT().Recv().Do(func() {
+				t.Logf("waiting close channel: %v", name)
 				xtest.WaitChannelClosed(t, streamClosed)
+				t.Logf("channel closed: %v", name)
 			}).Return(nil, errors.New("test stream closed")).MaxTimes(1)
 			return strm
 		}
 
-		strm2InitSent := make(empty.Chan)
-
-		strm2 := newStream("strm2", func() {
-			close(strm2InitSent)
-		})
+		strm2 := newStream("strm2")
 		strm2.EXPECT().Send(&rawtopicwriter.WriteRequest{
 			Messages: []rawtopicwriter.MessageData{
 				{SeqNo: 1},
 			},
 			Codec: rawtopiccommon.CodecRaw,
+		}).Do(func(_ *rawtopicwriter.WriteRequest) {
+			t.Logf("strm2 sent message and return retriable error")
 		}).Return(xerrors.Retryable(errors.New("retriable on strm2")))
 
-		strm3 := newStream("strm3", nil)
+		strm3 := newStream("strm3")
 		strm3.EXPECT().Send(&rawtopicwriter.WriteRequest{
 			Messages: []rawtopicwriter.MessageData{
 				{SeqNo: 1},
 			},
 			Codec: rawtopiccommon.CodecRaw,
+		}).Do(func(_ *rawtopicwriter.WriteRequest) {
+			t.Logf("strm3 sent message and return unretriable error")
 		}).Return(errors.New("strm3"))
 
 		connectsResult := []connectionAttemptContext{
@@ -488,17 +491,20 @@ func TestWriterImpl_Reconnect(t *testing.T) {
 			},
 		}
 
-		connectionAttempt := 0
+		var connectionAttempt xatomic.Int64
 		w.cfg.Connect = func(ctx context.Context) (RawTopicWriterStream, error) {
-			res := connectsResult[connectionAttempt]
-			connectionAttempt++
+			attemptIndex := int(connectionAttempt.Add(1)) - 1
+			t.Logf("connect with attempt index: %v", attemptIndex)
+			res := connectsResult[attemptIndex]
+
 			return res.stream, res.connectionError
 		}
 
 		connectionLoopStopped := make(empty.Chan)
 		go func() {
 			defer close(connectionLoopStopped)
 			w.connectionLoop(ctx)
+			t.Log("connection loop stopped")
 		}()
 
 		err := w.Write(ctx, newTestMessages(1))

tests/integration/basic_example_native_test.go

@@ -266,15 +266,11 @@ func TestBasicExampleNative(t *testing.T) { //nolint:gocyclo
 		ydb.WithSessionPoolSizeLimit(limit),
 		ydb.WithConnectionTTL(5*time.Second),
 		ydb.WithLogger(
-			log.Default(os.Stderr,
-				log.WithMinLevel(log.FromString(os.Getenv("YDB_LOG_SEVERITY_LEVEL"))),
-				log.WithColoring(),
-			),
+			newLoggerWithMinLevel(t, log.FromString(os.Getenv("YDB_LOG_SEVERITY_LEVEL"))),
 			trace.MatchDetails(`ydb\.(driver|table|discovery|retry|scheme).*`),
 		),
 		ydb.WithPanicCallback(func(e interface{}) {
-			_, _ = fmt.Fprintf(os.Stderr, "panic recovered:%v:\n%s", e, debug.Stack())
-			os.Exit(1)
+			t.Fatalf("panic recovered:%v:\n%s", e, debug.Stack())
 		}),
 		ydb.WithTraceTable(
 			*shutdownTrace.Compose(

tests/integration/connection_secure_test.go

@@ -25,7 +25,6 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3"
 	"github.com/ydb-platform/ydb-go-sdk/v3/config"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
-	"github.com/ydb-platform/ydb-go-sdk/v3/log"
 	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
@@ -53,9 +52,7 @@ func TestConnectionSecure(t *testing.T) {
 		ydb.WithConnectionTTL(time.Millisecond*10000),
 		ydb.WithMinTLSVersion(tls.VersionTLS10),
 		ydb.WithLogger(
-			log.Default(os.Stderr,
-				log.WithMinLevel(log.WARN),
-			),
+			newLogger(t),
 			trace.MatchDetails(`ydb\.(driver|discovery|retry|scheme).*`),
 		),
 	)

tests/integration/connection_test.go

@@ -84,9 +84,7 @@ func TestConnection(t *testing.T) {
 		ydb.WithConnectionTTL(time.Millisecond*10000),
 		ydb.WithMinTLSVersion(tls.VersionTLS10),
 		ydb.WithLogger(
-			log.Default(os.Stderr,
-				log.WithMinLevel(log.WARN),
-			),
+			newLoggerWithMinLevel(t, log.WARN),
 			trace.MatchDetails(`ydb\.(driver|discovery|retry|scheme).*`),
 		),
 		ydb.WithUserAgent(userAgent),

tests/integration/connection_with_compression_test.go

@@ -76,9 +76,7 @@ func TestConnectionWithCompression(t *testing.T) {
 		ydb.WithConnectionTTL(time.Millisecond*10000),
 		ydb.WithMinTLSVersion(tls.VersionTLS10),
 		ydb.WithLogger(
-			log.Default(os.Stderr,
-				log.WithMinLevel(log.WARN),
-			),
+			newLoggerWithMinLevel(t, log.WARN),
 			trace.MatchDetails(`ydb\.(driver|discovery|retry|scheme).*`),
 		),
 		ydb.WithUserAgent(userAgent),

tests/integration/database_sql_containers_test.go

@@ -7,7 +7,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"log"
 	"os"
 	"testing"
 
@@ -101,7 +100,7 @@ func TestDatabaseSqlContainers(t *testing.T) {
 				return rows.Scan(&testDatabaseSqlContainersVariantStruct{})
 			},
 			func() error {
-				return rows.Scan(&testDatabaseSqlContainersVariantTuple{})
+				return rows.Scan(&testDatabaseSqlContainersVariantTuple{test: t})
 			},
 			func() error {
 				var v int64
@@ -273,9 +272,10 @@ func (s *testDatabaseSqlContainersVariantStruct) Scan(res interface{}) error {
 }
 
 type testDatabaseSqlContainersVariantTuple struct {
-	a uint32
-	b string
-	c int64
+	a    uint32
+	b    string
+	c    int64
+	test testing.TB
 }
 
 func (s *testDatabaseSqlContainersVariantTuple) Scan(res interface{}) error {
@@ -305,8 +305,8 @@ func (s *testDatabaseSqlContainersVariantTuple) Scan(res interface{}) error {
 	return fmt.Errorf("type '%T' is not a `types.Value` type", res)
 }
 
-func (*testDatabaseSqlContainersVariantTuple) UnmarshalYDB(res types.RawValue) error {
-	log.Printf("T: %s", res.Type())
+func (s *testDatabaseSqlContainersVariantTuple) UnmarshalYDB(res types.RawValue) error {
+	s.test.Logf("T: %s", res.Type())
 	name, index := res.Variant()
 	var x interface{}
 	switch index {
@@ -317,7 +317,7 @@ func (*testDatabaseSqlContainersVariantTuple) UnmarshalYDB(res types.RawValue) e
 	case 2:
 		x = res.Int64()
 	}
-	log.Printf(
+	s.test.Logf(
 		"(tuple variant): %s %s %q %d = %v",
 		res.Path(), res.Type(), name, index, x,
 	)

tests/integration/discovery_test.go

@@ -63,10 +63,7 @@ func TestDiscovery(t *testing.T) {
 		ydb.WithConnectionTTL(time.Second*1),
 		ydb.WithMinTLSVersion(tls.VersionTLS10),
 		ydb.WithLogger(
-			log.Default(os.Stderr,
-				log.WithMinLevel(log.WARN),
-				log.WithColoring(),
-			),
+			newLoggerWithMinLevel(t, log.WARN),
 			trace.MatchDetails(`ydb\.(driver|discovery|repeater).*`),
 		),
 		ydb.WithUserAgent(userAgent),

tests/integration/helpers_test.go

@@ -7,8 +7,10 @@ import (
 	"bytes"
 	"context"
 	"database/sql"
+	"fmt"
 	"os"
 	"path"
+	"strings"
 	"testing"
 	"text/template"
 	"time"
@@ -17,6 +19,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
 	"github.com/ydb-platform/ydb-go-sdk/v3/log"
 	"github.com/ydb-platform/ydb-go-sdk/v3/sugar"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table"
@@ -90,9 +93,7 @@ func (scope *scopeT) Driver(opts ...ydb.Option) *ydb.Driver {
 			append(opts,
 				ydb.WithAccessTokenCredentials(token),
 				ydb.WithLogger(
-					log.Default(os.Stderr,
-						log.WithMinLevel(log.WARN),
-					),
+					scope.LoggerMinLevel(log.WARN),
 					trace.DetailsAll,
 				),
 			)...,
@@ -147,6 +148,18 @@ func (scope *scopeT) Folder() string {
 	}).(string)
 }
 
+func (scope *scopeT) Logger() *testLogger {
+	return scope.Cache(nil, nil, func() (res interface{}, err error) {
+		return newLogger(scope.t), nil
+	}).(*testLogger)
+}
+
+func (scope *scopeT) LoggerMinLevel(level log.Level) *testLogger {
+	return scope.Cache(level, nil, func() (res interface{}, err error) {
+		return newLoggerWithMinLevel(scope.t, level), nil
+	}).(*testLogger)
+}
+
 type tableNameParams struct {
 	tableName                string
 	createTableQueryTemplate string
@@ -224,3 +237,61 @@ func (scope *scopeT) TableName(opts ...func(t *tableNameParams)) string {
 func (scope *scopeT) TablePath(opts ...func(t *tableNameParams)) string {
 	return path.Join(scope.Folder(), scope.TableName(opts...))
 }
+
+// logger for tests
+type testLogger struct {
+	test     testing.TB
+	testName string
+	minLevel log.Level
+
+	m        xsync.Mutex
+	closed   bool
+	messages []string
+}
+
+func newLogger(t testing.TB) *testLogger {
+	return newLoggerWithMinLevel(t, 0)
+}
+
+func newLoggerWithMinLevel(t testing.TB, level log.Level) *testLogger {
+	logger := &testLogger{
+		test:     t,
+		testName: t.Name(),
+		minLevel: level,
+	}
+	t.Cleanup(logger.flush)
+	return logger
+}
+
+func (t *testLogger) Log(ctx context.Context, msg string, fields ...log.Field) {
+	t.test.Helper()
+	lvl := log.LevelFromContext(ctx)
+	if lvl < t.minLevel {
+		return
+	}
+
+	names := log.NamesFromContext(ctx)
+
+	loggerName := strings.Join(names, ".")
+	values := make(map[string]string)
+	for _, field := range fields {
+		values[field.Key()] = field.String()
+	}
+	timeString := time.Now().UTC().Format("15:04:05.999999999") // RFC3339Nano without date and timezone
+	message := fmt.Sprintf("%s: %s [%s] %s: %v (%v)", t.testName, timeString, lvl, loggerName, msg, values)
+	t.m.WithLock(func() {
+		if t.closed {
+			_, _ = fmt.Fprintf(os.Stderr, "\nFINISHED TEST %q:\n%s\n\n", t.testName, message)
+		} else {
+			t.messages = append(t.messages, message)
+		}
+	})
+}
+
+func (t *testLogger) flush() {
+	t.m.WithLock(func() {
+		t.closed = true
+		message := "\n" + strings.Join(t.messages, "\n")
+		t.test.Log(message)
+	})
+}

tests/integration/ratelimiter_test.go

@@ -37,9 +37,7 @@ func TestRatelimiter(t *testing.T) {
 		),
 		ydb.WithBalancer(balancers.SingleConn()),
 		ydb.WithLogger(
-			log.Default(os.Stderr,
-				log.WithMinLevel(log.WARN),
-			),
+			newLoggerWithMinLevel(t, log.WARN),
 			trace.MatchDetails(`ydb\.(driver|discovery|retry|ratelimiter|coordination).*`),
 		),
 	)

tests/integration/scripting_test.go

@@ -15,7 +15,6 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/balancers"
 	"github.com/ydb-platform/ydb-go-sdk/v3/config"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
-	"github.com/ydb-platform/ydb-go-sdk/v3/log"
 	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
 	"github.com/ydb-platform/ydb-go-sdk/v3/scripting"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table"
@@ -38,9 +37,7 @@ func TestScripting(t *testing.T) {
 		ydb.WithConnectionTTL(time.Millisecond*10000),
 		ydb.WithMinTLSVersion(tls.VersionTLS10),
 		ydb.WithLogger(
-			log.Default(os.Stderr,
-				log.WithMinLevel(log.TRACE),
-			),
+			newLogger(t),
 			trace.MatchDetails(`ydb\.(driver|discovery|retry|scheme).*`),
 		),
 		ydb.WithUserAgent("scripting"),