chore: dedup WaitForString() into testutils (#712)

Dedup the `WaitForString()` helper function that is used in `root_test` and db integration tests.
This commit is contained in:
AlexTalreja
2025-06-16 17:41:52 +00:00
committed by Alex Talreja
parent 8b0a18352e
commit a044c468d0
22 changed files with 105 additions and 148 deletions

View File

@@ -15,7 +15,6 @@
package cmd
import (
"bufio"
"bytes"
"context"
_ "embed"
@@ -992,67 +991,6 @@ func tmpFileWithCleanup(content []byte) (string, func(), error) {
return f.Name(), cleanup, err
}
// WaitForString is a helper function that waits for a string in the server log that matches the provided regex.
func WaitForString(ctx context.Context, re *regexp.Regexp, pr *io.PipeReader) (string, error) {
ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
defer cancel()
in := bufio.NewReader(pr)
type result struct {
s string
err error
}
output := make(chan result)
go func() {
defer close(output)
for {
select {
case <-ctx.Done():
// if the context is canceled, the orig thread will send back the error
// so we can just exit the goroutine here
return
default:
// otherwise read a line from the output
s, err := in.ReadString('\n')
if err != nil {
output <- result{err: err}
return
}
output <- result{s: s}
// if that last string matched, exit the goroutine
if re.MatchString(s) {
return
} else {
fmt.Printf("Actual logs: %s did not match pattern %s\n", s, re)
}
}
}
}()
// collect the output until the ctx is canceled, an error was hit,
// or match was found (which is indicated the channel is closed)
var sb strings.Builder
for {
select {
case <-ctx.Done():
// if ctx is done, return that error
return sb.String(), ctx.Err()
case o, ok := <-output:
if !ok {
// match was found!
return sb.String(), nil
}
if o.err != nil {
// error was found!
return sb.String(), o.err
}
sb.WriteString(o.s)
}
}
}
func TestSingleEdit(t *testing.T) {
ctx, cancelCtx := context.WithTimeout(context.Background(), time.Minute)
defer cancelCtx()
@@ -1077,9 +1015,9 @@ func TestSingleEdit(t *testing.T) {
go watchFile(ctx, fileToWatch)
begunWatchingFile := regexp.MustCompile(fmt.Sprintf("DEBUG \"Now watching tools file %s\"", fileToWatch))
_, err = WaitForString(ctx, begunWatchingFile, pr)
_, err = testutils.WaitForString(ctx, begunWatchingFile, pr)
if err != nil {
t.Fatalf("timeout or error waiting for watcher to start")
t.Fatalf("timeout or error waiting for watcher to start: %s", err)
}
err = os.WriteFile(fileToWatch, []byte("modification"), 0777)
@@ -1088,9 +1026,9 @@ func TestSingleEdit(t *testing.T) {
}
detectedFileChange := regexp.MustCompile(fmt.Sprintf("DEBUG \"WRITE event detected in tools file: %s", fileToWatch))
_, err = WaitForString(ctx, detectedFileChange, pr)
_, err = testutils.WaitForString(ctx, detectedFileChange, pr)
if err != nil {
t.Fatalf("timeout or error waiting for file to detect write %v", err)
t.Fatalf("timeout or error waiting for file to detect write: %s", err)
}
}

View File

@@ -15,9 +15,12 @@
package testutils
import (
"bufio"
"context"
"fmt"
"io"
"os"
"regexp"
"strings"
"github.com/googleapis/genai-toolbox/internal/log"
@@ -42,3 +45,63 @@ func ContextWithNewLogger() (context.Context, error) {
}
return util.WithLogger(ctx, logger), nil
}
// WaitForString waits until the server logs a single line that matches the provided regex.
// returns the output of whatever the server sent so far.
func WaitForString(ctx context.Context, re *regexp.Regexp, pr io.ReadCloser) (string, error) {
in := bufio.NewReader(pr)
ctx, cancel := context.WithCancel(ctx)
defer cancel()
// read lines in background, sending result of each read over a channel
// this allows us to use in.ReadString without blocking
type result struct {
s string
err error
}
output := make(chan result)
go func() {
defer close(output)
for {
select {
case <-ctx.Done():
// if the context is canceled, the orig thread will send back the error
// so we can just exit the goroutine here
return
default:
// otherwise read a line from the output
s, err := in.ReadString('\n')
if err != nil {
output <- result{err: err}
return
}
output <- result{s: s}
// if that last string matched, exit the goroutine
if re.MatchString(s) {
return
}
}
}
}()
// collect the output until the ctx is canceled, an error was hit,
// or match was found (which is indicated the channel is closed)
var sb strings.Builder
for {
select {
case <-ctx.Done():
// if ctx is done, return that error
return sb.String(), ctx.Err()
case o, ok := <-output:
if !ok {
// match was found!
return sb.String(), nil
}
if o.err != nil {
// error was found!
return sb.String(), o.err
}
sb.WriteString(o.s)
}
}
}

View File

@@ -28,6 +28,7 @@ import (
"time"
"github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
)
@@ -90,7 +91,7 @@ func TestAlloyDBAINLToolEndpoints(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)

View File

@@ -26,6 +26,7 @@ import (
"cloud.google.com/go/alloydbconn"
"github.com/google/uuid"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
"github.com/jackc/pgx/v5/pgxpool"
)
@@ -157,7 +158,7 @@ func TestAlloyDBPgToolEndpoints(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)

View File

@@ -29,6 +29,7 @@ import (
bigqueryapi "cloud.google.com/go/bigquery"
"github.com/google/uuid"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
"golang.org/x/oauth2/google"
"google.golang.org/api/googleapi"
@@ -124,7 +125,7 @@ func TestBigQueryToolEndpoints(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)

View File

@@ -29,6 +29,7 @@ import (
"cloud.google.com/go/bigtable"
"github.com/google/uuid"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/tests"
)
@@ -103,7 +104,7 @@ func TestBigtableToolEndpoints(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)

View File

@@ -29,6 +29,7 @@ import (
"cloud.google.com/go/cloudsqlconn"
"cloud.google.com/go/cloudsqlconn/sqlserver/mssql"
"github.com/google/uuid"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
)
@@ -151,7 +152,7 @@ func TestCloudSQLMSSQLToolEndpoints(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)

View File

@@ -28,6 +28,7 @@ import (
"cloud.google.com/go/cloudsqlconn"
"cloud.google.com/go/cloudsqlconn/mysql/mysql"
"github.com/google/uuid"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
)
@@ -138,7 +139,7 @@ func TestCloudSQLMySQLToolEndpoints(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)

View File

@@ -26,6 +26,7 @@ import (
"cloud.google.com/go/cloudsqlconn"
"github.com/google/uuid"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
"github.com/jackc/pgx/v5/pgxpool"
)
@@ -142,7 +143,7 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)

View File

@@ -25,6 +25,7 @@ import (
"github.com/couchbase/gocb/v2"
"github.com/google/uuid"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
)
@@ -128,7 +129,7 @@ func TestCouchbaseToolEndpoints(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)

View File

@@ -26,6 +26,7 @@ import (
"testing"
"time"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
)
@@ -78,7 +79,7 @@ func TestDgraphToolEndpoints(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)

View File

@@ -28,6 +28,7 @@ import (
"testing"
"time"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/tests"
)
@@ -254,7 +255,7 @@ func TestHttpToolEndpoints(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)

View File

@@ -26,6 +26,7 @@ import (
"time"
"github.com/google/uuid"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
)
@@ -124,7 +125,7 @@ func TestMSSQLToolEndpoints(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)

View File

@@ -25,6 +25,7 @@ import (
"time"
"github.com/google/uuid"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
)
@@ -115,7 +116,7 @@ func TestMySQLToolEndpoints(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)

View File

@@ -26,6 +26,7 @@ import (
"testing"
"time"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
)
@@ -87,7 +88,7 @@ func TestNeo4jToolEndpoints(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)

View File

@@ -25,6 +25,7 @@ import (
"time"
"github.com/google/uuid"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
"github.com/jackc/pgx/v5/pgxpool"
)
@@ -121,7 +122,7 @@ func TestPostgres(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)

View File

@@ -22,6 +22,7 @@ import (
"testing"
"time"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
"github.com/redis/go-redis/v9"
)
@@ -90,7 +91,7 @@ func TestRedisToolEndpoints(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)

View File

@@ -15,13 +15,10 @@
package tests
import (
"bufio"
"context"
"fmt"
"io"
"os"
"regexp"
"strings"
yaml "github.com/goccy/go-yaml"
@@ -133,63 +130,3 @@ func (c *CmdExec) Close() {
c.Close()
}
}
// WaitForString waits until the server logs a single line that matches the provided regex.
// returns the output of whatever the server sent so far.
func (c *CmdExec) WaitForString(ctx context.Context, re *regexp.Regexp) (string, error) {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
in := bufio.NewReader(c.Out)
// read lines in background, sending result of each read over a channel
// this allows us to use in.ReadString without blocking
type result struct {
s string
err error
}
output := make(chan result)
go func() {
defer close(output)
for {
select {
case <-ctx.Done():
// if the context is canceled, the orig thread will send back the error
// so we can just exit the goroutine here
return
default:
// otherwise read a line from the output
s, err := in.ReadString('\n')
if err != nil {
output <- result{err: err}
return
}
output <- result{s: s}
// if that last string matched, exit the goroutine
if re.MatchString(s) {
return
}
}
}
}()
// collect the output until the ctx is canceled, an error was hit,
// or match was found (which is indicated the channel is closed)
var sb strings.Builder
for {
select {
case <-ctx.Done():
// if ctx is done, return that error
return sb.String(), ctx.Err()
case o, ok := <-output:
if !ok {
// match was found!
return sb.String(), nil
}
if o.err != nil {
// error was found!
return sb.String(), o.err
}
sb.WriteString(o.s)
}
}
}

View File

@@ -26,6 +26,7 @@ import (
"time"
"cloud.google.com/go/cloudsqlconn"
"github.com/googleapis/genai-toolbox/internal/testutils"
)
// RunSourceConnection test for source connection
@@ -57,7 +58,7 @@ func RunSourceConnectionTest(t *testing.T, sourceConfig map[string]any, toolKind
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
return fmt.Errorf("toolbox didn't start successfully: %s", err)

View File

@@ -31,6 +31,7 @@ import (
database "cloud.google.com/go/spanner/admin/database/apiv1"
"cloud.google.com/go/spanner/admin/database/apiv1/databasepb"
"github.com/google/uuid"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/tests"
)
@@ -141,7 +142,7 @@ func TestSpannerToolEndpoints(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)

View File

@@ -25,6 +25,7 @@ import (
"time"
"github.com/google/uuid"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
)
@@ -143,7 +144,7 @@ func TestSQLiteToolEndpoint(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)

View File

@@ -22,6 +22,7 @@ import (
"testing"
"time"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
"github.com/valkey-io/valkey-go"
)
@@ -93,7 +94,7 @@ func TestValkeyToolEndpoints(t *testing.T) {
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)