gcsbuf/writer_test.go

package buf

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

	"cloud.google.com/go/storage"
	"github.com/fsouza/fake-gcs-server/fakestorage"
)

func setup(t *testing.T) *storage.Client {
	t.Helper()
	server := fakestorage.NewServer(nil)
	server.CreateBucketWithOpts(fakestorage.CreateBucketOpts{Name: "test"})
	t.Cleanup(server.Stop)
	return server.Client()
}

// readObject reads the full content of a GCS object as a string.
func readObject(t *testing.T, client *storage.Client, bucket, object string) string {
	t.Helper()
	rc, err := client.Bucket(bucket).Object(object).NewReader(context.Background())
	if err != nil {
		t.Fatalf("reading object %s/%s: %v", bucket, object, err)
	}
	defer rc.Close()
	data, err := io.ReadAll(rc)
	if err != nil {
		t.Fatalf("reading object %s/%s: %v", bucket, object, err)
	}
	return string(data)
}

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

	w := NewWriter(ctx, WriterOptions{
		Client:        client,
		Bucket:        "test",
		Object:        "obj",
		FlushInterval: time.Hour, // disable periodic flush
	})

	data := []byte("hello world")
	n, err := w.Write(data)
	if err != nil {
		t.Fatal(err)
	}
	if n != len(data) {
		t.Fatalf("wrote %d, want %d", n, len(data))
	}

	if err := w.Flush(); err != nil {
		t.Fatal(err)
	}

	// Verify object content.
	got := readObject(t, client, "test", "obj")
	if got != "hello world" {
		t.Fatalf("got %q, want %q", got, "hello world")
	}
}

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

	w := NewWriter(ctx, WriterOptions{
		Client:        client,
		Bucket:        "test",
		Object:        "obj",
		FlushInterval: time.Hour,
	})

	w.Write([]byte("aaa"))
	if err := w.Flush(); err != nil {
		t.Fatal(err)
	}
	w.Write([]byte("bbb"))
	if err := w.Flush(); err != nil {
		t.Fatal(err)
	}
	w.Write([]byte("ccc"))
	if err := w.Flush(); err != nil {
		t.Fatal(err)
	}

	got := readObject(t, client, "test", "obj")
	if got != "aaabbbccc" {
		t.Fatalf("got %q, want %q", got, "aaabbbccc")
	}
}

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

	w := NewWriter(ctx, WriterOptions{
		Client:        client,
		Bucket:        "test",
		Object:        "obj",
		FlushInterval: time.Hour,
		FlushSize:     10, // small threshold
	})

	// Write more than flushSize; should auto-flush.
	data := []byte("0123456789extra")
	if _, err := w.Write(data); err != nil {
		t.Fatal(err)
	}

	// Object should exist now (auto-flushed).
	got := readObject(t, client, "test", "obj")
	if got != string(data) {
		t.Fatalf("got %q, want %q", got, string(data))
	}
}

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

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

	w.Write([]byte("periodic"))

	// Wait for periodic flush.
	time.Sleep(200 * time.Millisecond)

	got := readObject(t, client, "test", "obj")
	if got != "periodic" {
		t.Fatalf("got %q, want %q", got, "periodic")
	}

	if err := w.Close(); err != nil {
		t.Fatal(err)
	}
}

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

	w := NewWriter(ctx, WriterOptions{
		Client:        client,
		Bucket:        "test",
		Object:        "obj",
		FlushInterval: time.Hour,
	})

	w.Write([]byte("data"))
	if err := w.Close(); err != nil {
		t.Fatal(err)
	}

	attrs, err := client.Bucket("test").Object("obj").Attrs(ctx)
	if err != nil {
		t.Fatal(err)
	}
	if attrs.Metadata["done"] != "true" {
		t.Fatalf("metadata = %v, want done=true", attrs.Metadata)
	}
}

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

	w := NewWriter(ctx, WriterOptions{
		Client:        client,
		Bucket:        "test",
		Object:        "obj",
		FlushInterval: time.Hour,
	})

	w.Write([]byte("data"))
	w.Close()

	_, err := w.Write([]byte("more"))
	if err == nil {
		t.Fatal("expected error writing to closed writer")
	}
}

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

	w := NewWriter(ctx, WriterOptions{
		Client:        client,
		Bucket:        "test",
		Object:        "obj",
		FlushInterval: time.Hour,
	})

	// Close without writing anything should not error.
	if err := w.Close(); err != nil {
		t.Fatal(err)
	}
}

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

	// External process creates the object first.
	extw := client.Bucket("test").Object("obj").NewWriter(ctx)
	extw.Write([]byte("external"))
	if err := extw.Close(); err != nil {
		t.Fatalf("external write: %v", err)
	}

	// Our writer tries to create the same object — DoesNotExist should fail.
	w := NewWriter(ctx, WriterOptions{
		Client:        client,
		Bucket:        "test",
		Object:        "obj",
		FlushInterval: time.Hour,
	})

	w.Write([]byte("hello"))
	err := w.Flush()
	if err == nil {
		t.Fatal("expected error from flush when object already exists, got nil")
	}
	t.Logf("got expected error: %v", err)
}

func TestWriterGenerationMismatchDetected(t *testing.T) {
	fake := newFakeGCS()
	client := fake.client(t)
	ctx := context.Background()

	w := NewWriter(ctx, WriterOptions{
		Client:        client,
		Bucket:        "test",
		Object:        "obj",
		FlushInterval: time.Hour,
	})

	// First write+flush: creates the object.
	w.Write([]byte("hello"))
	if err := w.Flush(); err != nil {
		t.Fatalf("first flush: %v", err)
	}

	// Externally overwrite the object (simulating another writer).
	extw := client.Bucket("test").Object("obj").NewWriter(ctx)
	extw.Write([]byte("external"))
	if err := extw.Close(); err != nil {
		t.Fatalf("external write: %v", err)
	}

	// The writer's tracked generation is now stale.
	// Next flush should fail with a precondition error (412).
	w.Write([]byte(" world"))
	err := w.Flush()
	if err == nil {
		t.Fatal("expected error from flush after external overwrite, got nil")
	}
	t.Logf("got expected error: %v", err)
}

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

	w := NewWriter(ctx, WriterOptions{
		Client:        client,
		Bucket:        "test",
		Object:        "obj",
		FlushInterval: time.Hour,
		FlushSize:     1024 * 1024, // large so we control flushes manually
	})

	// Do many flushes to trigger component count reset.
	var want strings.Builder
	for i := range 35 {
		chunk := strings.Repeat(string(rune('a'+i%26)), 100)
		want.WriteString(chunk)
		w.Write([]byte(chunk))
		if err := w.Flush(); err != nil {
			t.Fatalf("flush %d: %v", i, err)
		}
	}

	if err := w.Close(); err != nil {
		t.Fatal(err)
	}

	got := readObject(t, client, "test", "obj")
	if got != want.String() {
		t.Fatalf("content mismatch after %d flushes: got len=%d, want len=%d", 35, len(got), want.Len())
	}
}