logarchiver/internal/consumer/runner_test.go

package consumer

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

	"git.unkin.net/unkin/logarchiver/internal/batcher"
	"github.com/nats-io/nats.go"
	"github.com/nats-io/nats.go/jetstream"
)

// fakeMsg is a minimal jetstream.Msg recording ack/nak calls.
type fakeMsg struct {
	subject string
	data    []byte
	mu      sync.Mutex
	acked   bool
	naked   bool
}

func (m *fakeMsg) Metadata() (*jetstream.MsgMetadata, error) { return &jetstream.MsgMetadata{}, nil }
func (m *fakeMsg) Data() []byte                              { return m.data }
func (m *fakeMsg) Headers() nats.Header                      { return nil }
func (m *fakeMsg) Subject() string                           { return m.subject }
func (m *fakeMsg) Reply() string                             { return "" }
func (m *fakeMsg) Ack() error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.acked = true
	return nil
}
func (m *fakeMsg) DoubleAck(context.Context) error { return nil }
func (m *fakeMsg) Nak() error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.naked = true
	return nil
}
func (m *fakeMsg) NakWithDelay(time.Duration) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.naked = true
	return nil
}
func (m *fakeMsg) InProgress() error           { return nil }
func (m *fakeMsg) Term() error                 { return nil }
func (m *fakeMsg) TermWithReason(string) error { return nil }

func (m *fakeMsg) isAcked() bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.acked
}
func (m *fakeMsg) isNaked() bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.naked
}

// fakePersister records calls and can be made to fail.
type fakePersister struct {
	fail   bool
	called int
}

func (p *fakePersister) Store(_ context.Context, b *batcher.Batch) (StoreResult, error) {
	p.called++
	if p.fail {
		return StoreResult{}, errors.New("boom")
	}
	return StoreResult{ObjectKey: "k", Events: len(b.Items)}, nil
}

func batchWith(msgs ...*fakeMsg) *batcher.Batch {
	b := &batcher.Batch{Subject: "s"}
	for _, m := range msgs {
		b.Items = append(b.Items, batcher.Item{Subject: "s", Raw: m.data, Ack: jetstream.Msg(m)})
	}
	return b
}

// TestFlushAcksOnlyAfterPersist is the core correctness test: messages are acked
// exactly when Store succeeds, and Nak'd (never acked) when it fails.
func TestFlushAcksOnSuccess(t *testing.T) {
	p := &fakePersister{}
	r := NewRunner(Options{Persister: p})
	m1 := &fakeMsg{subject: "s", data: []byte(`{"host":"h"}`)}
	m2 := &fakeMsg{subject: "s", data: []byte(`{"host":"h2"}`)}

	r.flush(context.Background(), batchWith(m1, m2), "test")

	if p.called != 1 {
		t.Fatalf("Store called %d times, want 1", p.called)
	}
	if !m1.isAcked() || !m2.isAcked() {
		t.Errorf("messages should be acked after successful persist")
	}
	if m1.isNaked() || m2.isNaked() {
		t.Errorf("messages must not be naked on success")
	}
}

func TestFlushNaksOnFailure(t *testing.T) {
	p := &fakePersister{fail: true}
	r := NewRunner(Options{Persister: p})
	m1 := &fakeMsg{subject: "s", data: []byte(`{"host":"h"}`)}

	r.flush(context.Background(), batchWith(m1), "test")

	if m1.isAcked() {
		t.Errorf("message must NOT be acked when persist fails")
	}
	if !m1.isNaked() {
		t.Errorf("message should be naked so JetStream redelivers")
	}
}

func TestFlushEmptyBatchNoop(t *testing.T) {
	p := &fakePersister{}
	r := NewRunner(Options{Persister: p})
	r.flush(context.Background(), &batcher.Batch{Subject: "s"}, "test")
	if p.called != 0 {
		t.Errorf("empty batch should not call Store")
	}
}

// TestRouteAndFlushIntegration wires a real batcher: adding enough messages to
// fill the batch triggers a full flush that persists and acks exactly those.
func TestRouteFlushViaBatcher(t *testing.T) {
	p := &fakePersister{}
	bat := batcher.New(batcher.Limits{MaxEvents: 2})
	r := NewRunner(Options{Persister: p, Batcher: bat})

	m1 := &fakeMsg{subject: "s", data: []byte(`{"host":"a"}`)}
	m2 := &fakeMsg{subject: "s", data: []byte(`{"host":"b"}`)}
	r.route(context.Background(), m1)
	if m1.isAcked() {
		t.Errorf("first message should not be acked before batch fills")
	}
	r.route(context.Background(), m2) // fills batch -> flush

	if p.called != 1 {
		t.Fatalf("Store called %d times, want 1 after fill", p.called)
	}
	if !m1.isAcked() || !m2.isAcked() {
		t.Errorf("both messages should be acked after the full-batch flush")
	}
}