From 045b8f8cab70e96a909276c728562effb96c5117 Mon Sep 17 00:00:00 2001
From: sashayakovtseva <sashayakovtseva@gmail.com>
Date: Tue, 6 Feb 2024 15:26:01 +0800
Subject: [PATCH] Do not check local YDB with exposed port

Signed-off-by: sashayakovtseva <sashayakovtseva@gmail.com>
---
 internal/testserver/datastore/ydb.go | 53 +++++++++++++++++-----------
 1 file changed, 32 insertions(+), 21 deletions(-)

diff --git a/internal/testserver/datastore/ydb.go b/internal/testserver/datastore/ydb.go
index 714b059919..4244d011ed 100644
--- a/internal/testserver/datastore/ydb.go
+++ b/internal/testserver/datastore/ydb.go
@@ -12,16 +12,16 @@ import (
 	"github.com/google/uuid"
 	"github.com/ory/dockertest/v3"
 	"github.com/stretchr/testify/require"
-	"github.com/ydb-platform/ydb-go-sdk/v3"
 
 	ydbDatastore "github.com/authzed/spicedb/internal/datastore/ydb"
 	ydbMigrations "github.com/authzed/spicedb/internal/datastore/ydb/migrations"
+
 	"github.com/authzed/spicedb/pkg/datastore"
 	"github.com/authzed/spicedb/pkg/migrate"
 )
 
 const (
-	ydbTestVersionTag  = "23.3.17"
+	ydbTestVersionTag  = "nightly"
 	ydbDefaultDatabase = "local"
 	ydbGRPCPort        = 2136
 )
@@ -51,40 +51,51 @@ func (r ydbTester) NewDatabase(t testing.TB) string {
 	// create a new container with default /local database instead.
 
 	containerName := fmt.Sprintf("ydb-%s", uuid.New().String())
+	hostname := "localhost"
+	if r.bridgeNetworkName != "" {
+		hostname = containerName
+	}
+
 	resource, err := r.pool.RunWithOptions(&dockertest.RunOptions{
 		Name:       containerName,
-		Hostname:   "localhost",
-		Repository: "cr.yandex/yc/yandex-docker-local-ydb",
+		Hostname:   hostname,
+		Repository: "ghcr.io/ydb-platform/local-ydb",
 		Tag:        ydbTestVersionTag,
-		Env:        []string{"YDB_USE_IN_MEMORY_PDISKS=true"},
-		NetworkID:  r.bridgeNetworkName,
+		Env: []string{
+			"YDB_USE_IN_MEMORY_PDISKS=true",
+			"YDB_FEATURE_FLAGS=enable_not_null_data_columns",
+		},
+		NetworkID: r.bridgeNetworkName,
 	})
 	require.NoError(t, err)
 	t.Cleanup(func() { require.NoError(t, r.pool.Purge(resource)) })
 
-	hostname := "localhost"
-	port := resource.GetPort(fmt.Sprintf("%d/tcp", ydbGRPCPort))
-	if r.bridgeNetworkName != "" {
-		hostname = containerName
-		port = strconv.FormatInt(ydbGRPCPort, 10)
-	}
-
-	dsn := fmt.Sprintf("grpc://%s:%s/%s", hostname, port, ydbDefaultDatabase)
 	require.NoError(t, r.pool.Retry(func() error { // await container is ready
-		ctx, cancel := context.WithTimeout(context.Background(), dockerBootTimeout)
-		defer cancel()
-
-		driver, err := ydb.Open(ctx, dsn)
+		code, err := resource.Exec([]string{
+			"/ydb",
+			"-e",
+			"grpc://localhost:2136",
+			"-d",
+			"/local",
+			"scheme",
+			"ls",
+		}, dockertest.ExecOptions{})
 		if err != nil {
 			return err
 		}
-
-		if _, err := driver.Scheme().ListDirectory(ctx, ydbDefaultDatabase); err != nil {
-			return err
+		if code != 0 {
+			return fmt.Errorf("exited with %d", code)
 		}
 
 		return nil
 	}))
+
+	port := resource.GetPort(fmt.Sprintf("%d/tcp", ydbGRPCPort))
+	if r.bridgeNetworkName != "" {
+		port = strconv.FormatInt(ydbGRPCPort, 10)
+	}
+
+	dsn := fmt.Sprintf("grpc://%s:%s/%s", hostname, port, ydbDefaultDatabase)
 	return dsn
 }