OpenFlare/backend/pkg/batchwriter/writer_test.go

// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0

package batchwriter

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type testEvent struct {
	ID   int
	Data string
}

func testConfig() Config {
	return Config{
		Name:          "test-writer",
		QueueSize:     100,
		MaxBatchSize:  5,
		FlushInterval: 20 * time.Millisecond,
	}
}

func TestWriter_BatchSizeFlush(t *testing.T) {
	var (
		mu      sync.Mutex
		batches [][]testEvent
		flushWg sync.WaitGroup
	)
	flushWg.Add(1)

	cfg := testConfig()
	cfg.FlushInterval = time.Hour

	w, err := New(cfg, func(_ context.Context, items []testEvent) error {
		mu.Lock()
		defer mu.Unlock()
		batches = append(batches, items)
		if len(items) == 5 {
			flushWg.Done()
		}
		return nil
	})
	require.NoError(t, err)

	w.Start(context.Background())
	defer func() { _ = w.Stop(context.Background()) }()

	for i := 1; i <= 5; i++ {
		ok := w.TryEnqueue(testEvent{ID: i, Data: "payload"})
		assert.True(t, ok)
	}

	done := make(chan struct{})
	go func() {
		flushWg.Wait()
		close(done)
	}()

	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for batch flush")
	}

	mu.Lock()
	defer mu.Unlock()
	require.Len(t, batches, 1)
	assert.Len(t, batches[0], 5)
	for i, item := range batches[0] {
		assert.Equal(t, i+1, item.ID)
	}
}

func TestWriter_IntervalFlush(t *testing.T) {
	var (
		mu      sync.Mutex
		flushed []testEvent
		done    = make(chan struct{})
	)

	cfg := testConfig()
	cfg.MaxBatchSize = 100
	cfg.FlushInterval = 30 * time.Millisecond

	w, err := New(cfg, func(_ context.Context, items []testEvent) error {
		mu.Lock()
		defer mu.Unlock()
		flushed = append(flushed, items...)
		if len(flushed) == 2 {
			select {
			case <-done:
			default:
				close(done)
			}
		}
		return nil
	})
	require.NoError(t, err)

	w.Start(context.Background())
	defer func() { _ = w.Stop(context.Background()) }()

	assert.True(t, w.TryEnqueue(testEvent{ID: 1}))
	assert.True(t, w.TryEnqueue(testEvent{ID: 2}))

	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for interval flush")
	}

	mu.Lock()
	defer mu.Unlock()
	assert.Len(t, flushed, 2)
}

func TestWriter_MinBatchSizeThreshold(t *testing.T) {
	var (
		mu      sync.Mutex
		flushed []testEvent
		done    = make(chan struct{})
	)

	cfg := testConfig()
	cfg.MaxBatchSize = 100
	cfg.MinBatchSize = 3
	cfg.FlushInterval = 20 * time.Millisecond
	cfg.MaxFlushWait = 60 * time.Millisecond

	w, err := New(cfg, func(_ context.Context, items []testEvent) error {
		mu.Lock()
		defer mu.Unlock()
		flushed = append(flushed, items...)
		if len(flushed) == 2 {
			select {
			case <-done:
			default:
				close(done)
			}
		}
		return nil
	})
	require.NoError(t, err)

	w.Start(context.Background())
	defer func() { _ = w.Stop(context.Background()) }()

	assert.True(t, w.TryEnqueue(testEvent{ID: 1}))
	assert.True(t, w.TryEnqueue(testEvent{ID: 2}))

	time.Sleep(30 * time.Millisecond)
	mu.Lock()
	assert.Empty(t, flushed, "items should wait until MinBatchSize or MaxFlushWait")
	mu.Unlock()

	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for forced max wait flush")
	}

	mu.Lock()
	defer mu.Unlock()
	assert.Len(t, flushed, 2)
}

func TestWriter_StopDrainsRemaining(t *testing.T) {
	var (
		mu      sync.Mutex
		flushed []testEvent
	)

	cfg := testConfig()
	cfg.FlushInterval = time.Hour

	w, err := New(cfg, func(_ context.Context, items []testEvent) error {
		mu.Lock()
		defer mu.Unlock()
		flushed = append(flushed, items...)
		return nil
	})
	require.NoError(t, err)

	w.Start(context.Background())
	for i := 1; i <= 3; i++ {
		assert.True(t, w.TryEnqueue(testEvent{ID: i}))
	}

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	require.NoError(t, w.Stop(ctx))

	assert.False(t, w.Running())
	mu.Lock()
	defer mu.Unlock()
	assert.Len(t, flushed, 3)
}

func TestWriter_DropWhenFull(t *testing.T) {
	var (
		dropped atomic.Int64
		blockCh = make(chan struct{})
	)

	cfg := Config{
		QueueSize:     2,
		MaxBatchSize:  1,
		FlushInterval: time.Hour,
	}

	entered := make(chan struct{})
	w, err := New(cfg, func(_ context.Context, _ []testEvent) error {
		select {
		case entered <- struct{}{}:
		default:
		}
		<-blockCh
		return nil
	}, WithDropHandler(func(_ testEvent) {
		dropped.Add(1)
	}))
	require.NoError(t, err)

	w.Start(context.Background())
	defer func() {
		close(blockCh)
		_ = w.Stop(context.Background())
	}()

	// 1. 推入 1 个 item 触发 flush 并阻塞在 blockCh
	w.ch <- testEvent{ID: 1}
	<-entered

	// 2. 此时 worker 阻塞，填满 channel
	w.ch <- testEvent{ID: 2}
	w.ch <- testEvent{ID: 3}

	assert.False(t, w.TryEnqueue(testEvent{ID: 4}))
	assert.Equal(t, int64(1), dropped.Load())
	assert.Equal(t, int64(1), w.Stats().Drops)
}

func TestWriter_FlushErrorCallback(t *testing.T) {
	var (
		called   atomic.Bool
		flushErr = errors.New("clickhouse write timeout")
		done     = make(chan struct{})
	)

	cfg := testConfig()
	cfg.MaxBatchSize = 1
	cfg.FlushInterval = time.Hour

	w, err := New(cfg, func(_ context.Context, _ []testEvent) error {
		return flushErr
	}, WithFlushErrorHandler(func(_ context.Context, items []testEvent, err error) {
		called.Store(true)
		assert.Equal(t, flushErr, err)
		assert.Len(t, items, 1)
		close(done)
	}))
	require.NoError(t, err)

	w.Start(context.Background())
	defer func() { _ = w.Stop(context.Background()) }()

	assert.True(t, w.TryEnqueue(testEvent{ID: 1}))

	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for error callback")
	}

	assert.True(t, called.Load())
	assert.Equal(t, int64(1), w.Stats().FlushErrors)
}

func TestWriter_ValidateConfig(t *testing.T) {
	tests := []struct {
		name    string
		cfg     Config
		wantErr bool
	}{
		{"valid", DefaultConfig(), false},
		{"zero queue", Config{QueueSize: 0, MaxBatchSize: 10, FlushInterval: time.Second}, true},
		{"zero max batch", Config{QueueSize: 10, MaxBatchSize: 0, FlushInterval: time.Second}, true},
		{"negative min batch", Config{QueueSize: 10, MaxBatchSize: 10, MinBatchSize: -1, FlushInterval: time.Second}, true},
		{"zero flush interval", Config{QueueSize: 10, MaxBatchSize: 10, FlushInterval: 0}, true},
		{"negative max flush wait", Config{QueueSize: 10, MaxBatchSize: 10, FlushInterval: time.Second, MaxFlushWait: -1}, true},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			_, err := New(tt.cfg, func(_ context.Context, _ []testEvent) error { return nil })
			if tt.wantErr {
				assert.Error(t, err)
			} else {
				assert.NoError(t, err)
			}
		})
	}
}

func TestWriter_NilFlushFunc(t *testing.T) {
	_, err := New[testEvent](DefaultConfig(), nil)
	assert.ErrorIs(t, err, errNilFlushFunc)
}