gcsbuf/reader_test.go

package buf

import (
	"bytes"
	"context"
	"io"
	"strings"
	"sync"
	"testing"
	"time"

	"cloud.google.com/go/storage"
)

// writeObject writes content to a GCS object with optional metadata.
func writeObject(t *testing.T, client *storage.Client, bucket, object, content string, metadata map[string]string) {
	t.Helper()
	ctx := context.Background()
	w := client.Bucket(bucket).Object(object).NewWriter(ctx)
	if metadata != nil {
		w.Metadata = metadata
	}
	if _, err := io.Copy(w, bytes.NewReader([]byte(content))); err != nil {
		t.Fatalf("writing object %s/%s: %v", bucket, object, err)
	}
	if err := w.Close(); err != nil {
		t.Fatalf("closing object writer %s/%s: %v", bucket, object, err)
	}
}

func TestReaderFromCompleteObject(t *testing.T) {
	client := setup(t)
	ctx := context.Background()

	// Pre-populate an object with done metadata.
	writeObject(t, client, "test", "complete", "hello reader", map[string]string{"done": "true"})

	r := NewReader(ctx, ReaderOptions{
		Client:       client,
		Bucket:       "test",
		Object:       "complete",
		PollInterval: 10 * time.Millisecond,
	})

	got, err := io.ReadAll(r)
	if err != nil {
		t.Fatal(err)
	}
	if string(got) != "hello reader" {
		t.Fatalf("got %q, want %q", string(got), "hello reader")
	}
}

func TestReaderPollsForObjectCreation(t *testing.T) {
	client := setup(t)
	ctx := context.Background()

	r := NewReader(ctx, ReaderOptions{
		Client:       client,
		Bucket:       "test",
		Object:       "delayed",
		PollInterval: 10 * time.Millisecond,
	})

	// Create the object after a short delay.
	go func() {
		time.Sleep(50 * time.Millisecond)
		writeObject(t, client, "test", "delayed", "delayed data", map[string]string{"done": "true"})
	}()

	got, err := io.ReadAll(r)
	if err != nil {
		t.Fatal(err)
	}
	if string(got) != "delayed data" {
		t.Fatalf("got %q, want %q", string(got), "delayed data")
	}
}

func TestReaderPollsForMoreData(t *testing.T) {
	client := setup(t)
	ctx := context.Background()

	// Create initial object without done.
	writeObject(t, client, "test", "growing", "part1", nil)

	r := NewReader(ctx, ReaderOptions{
		Client:       client,
		Bucket:       "test",
		Object:       "growing",
		PollInterval: 10 * time.Millisecond,
	})

	// Read first chunk.
	buf := make([]byte, 100)
	n, err := r.Read(buf)
	if err != nil {
		t.Fatal(err)
	}
	if string(buf[:n]) != "part1" {
		t.Fatalf("got %q, want %q", string(buf[:n]), "part1")
	}

	// Overwrite with more data and set done.
	go func() {
		time.Sleep(50 * time.Millisecond)
		writeObject(t, client, "test", "growing", "part1part2", map[string]string{"done": "true"})
	}()

	rest, err := io.ReadAll(r)
	if err != nil {
		t.Fatal(err)
	}
	if string(rest) != "part2" {
		t.Fatalf("got %q, want %q", string(rest), "part2")
	}
}

func TestReaderContextCancellation(t *testing.T) {
	client := setup(t)
	ctx, cancel := context.WithCancel(context.Background())

	r := NewReader(ctx, ReaderOptions{
		Client:       client,
		Bucket:       "test",
		Object:       "nonexistent",
		PollInterval: 10 * time.Millisecond,
	})

	go func() {
		time.Sleep(50 * time.Millisecond)
		cancel()
	}()

	_, err := r.Read(make([]byte, 100))
	if err == nil {
		t.Fatal("expected error from cancelled context")
	}
}

func TestIntegrationWriterReader(t *testing.T) {
	client := setup(t)
	ctx := context.Background()

	w := NewWriter(ctx, WriterOptions{
		Client:        client,
		Bucket:        "test",
		Object:        "stream",
		FlushInterval: 50 * time.Millisecond,
		FlushSize:     1024 * 1024,
	})

	r := NewReader(ctx, ReaderOptions{
		Client:       client,
		Bucket:       "test",
		Object:       "stream",
		PollInterval: 10 * time.Millisecond,
	})

	// Write data in chunks.
	chunks := []string{"hello ", "world ", "from ", "gcsbuf"}
	var wg sync.WaitGroup
	wg.Go(func() {
		for _, chunk := range chunks {
			if _, err := w.Write([]byte(chunk)); err != nil {
				t.Errorf("write: %v", err)
				return
			}
			time.Sleep(30 * time.Millisecond)
		}
		if err := w.Close(); err != nil {
			t.Errorf("close: %v", err)
		}
	})

	got, err := io.ReadAll(r)
	if err != nil {
		t.Fatal(err)
	}

	wg.Wait()

	want := "hello world from gcsbuf"
	if string(got) != want {
		t.Fatalf("got %q, want %q", string(got), want)
	}
}

func TestIntegrationMultipleFlushCyclesWithReader(t *testing.T) {
	client := setup(t)
	ctx := context.Background()

	w := NewWriter(ctx, WriterOptions{
		Client:        client,
		Bucket:        "test",
		Object:        "multi",
		FlushInterval: time.Hour,
		FlushSize:     10, // flush every 10 bytes
	})

	r := NewReader(ctx, ReaderOptions{
		Client:       client,
		Bucket:       "test",
		Object:       "multi",
		PollInterval: 10 * time.Millisecond,
	})

	// Write enough data to trigger multiple auto-flushes.
	var wg sync.WaitGroup
	wg.Go(func() {
		for i := range 5 {
			chunk := strings.Repeat("x", 15) // > flushSize, triggers auto-flush each time
			if _, err := w.Write([]byte(chunk)); err != nil {
				t.Errorf("write %d: %v", i, err)
				return
			}
		}
		if err := w.Close(); err != nil {
			t.Errorf("close: %v", err)
		}
	})

	got, err := io.ReadAll(r)
	if err != nil {
		t.Fatal(err)
	}

	wg.Wait()

	want := strings.Repeat("x", 75)
	if string(got) != want {
		t.Fatalf("got len=%d, want len=%d", len(got), len(want))
	}
}