lingniu-vehicle-ingest/go/vehicle-gateway/cmd/nats-fast-writer/main_test.go

package main

import (
	"context"
	"errors"
	"strings"
	"sync"
	"testing"

	"github.com/nats-io/nats.go"

	"lingniu-vehicle-ingest/go/vehicle-gateway/internal/envelope"
	"lingniu-vehicle-ingest/go/vehicle-gateway/internal/metrics"
)

func TestProcessFastMessageWritesTDengineAndRedisBeforeAck(t *testing.T) {
	env := envelope.FrameEnvelope{Protocol: envelope.ProtocolGB32960, VIN: "VIN001", EventID: "evt-1"}
	payload, err := env.MarshalJSONBytes()
	if err != nil {
		t.Fatal(err)
	}
	appender := &recordingFastAppender{}
	updater := &recordingFastUpdater{}
	ackCount := 0
	msg := &fastMessage{data: payload, ack: func() error {
		ackCount++
		return nil
	}}

	if err := processFastMessage(context.Background(), nil, appender, updater, msg); err != nil {
		t.Fatalf("processFastMessage() error = %v", err)
	}
	if appender.count != 1 || updater.count != 1 || ackCount != 1 {
		t.Fatalf("appends=%d updates=%d acks=%d", appender.count, updater.count, ackCount)
	}
}

func TestProcessFastMessageDoesNotAckWhenRedisFails(t *testing.T) {
	env := envelope.FrameEnvelope{Protocol: envelope.ProtocolJT808, VIN: "VIN001", EventID: "evt-2"}
	payload, err := env.MarshalJSONBytes()
	if err != nil {
		t.Fatal(err)
	}
	appender := &recordingFastAppender{}
	updater := &recordingFastUpdater{err: errors.New("redis down")}
	ackCount := 0
	msg := &fastMessage{data: payload, ack: func() error {
		ackCount++
		return nil
	}}

	if err := processFastMessage(context.Background(), nil, appender, updater, msg); err == nil {
		t.Fatal("processFastMessage() error = nil, want redis failure")
	}
	if ackCount != 0 {
		t.Fatalf("ack count = %d, want 0", ackCount)
	}
}

func TestProcessFastMessageRecordsStageDurationMetrics(t *testing.T) {
	env := envelope.FrameEnvelope{Protocol: envelope.ProtocolJT808, VIN: "VIN001", EventID: "evt-3"}
	payload, err := env.MarshalJSONBytes()
	if err != nil {
		t.Fatal(err)
	}
	registry := metrics.NewRegistry()
	msg := &fastMessage{subject: "vehicle.raw.go.jt808.v1", data: payload, ack: func() error { return nil }}

	if err := processFastMessage(context.Background(), registry, &recordingFastAppender{}, &recordingFastUpdater{}, msg); err != nil {
		t.Fatalf("processFastMessage() error = %v", err)
	}

	text := registry.Render()
	for _, want := range []string{
		`vehicle_fast_writer_stage_duration_ms_histogram_bucket{le="+Inf",stage="tdengine",status="ok",subject="vehicle.raw.go.jt808.v1"} 1`,
		`vehicle_fast_writer_stage_duration_ms_histogram_count{stage="tdengine",status="ok",subject="vehicle.raw.go.jt808.v1"} 1`,
		`vehicle_fast_writer_stage_duration_ms_histogram_bucket{le="+Inf",stage="redis",status="ok",subject="vehicle.raw.go.jt808.v1"} 1`,
		`vehicle_fast_writer_stage_duration_ms_histogram_bucket{le="+Inf",stage="ack",status="ok",subject="vehicle.raw.go.jt808.v1"} 1`,
	} {
		if !strings.Contains(text, want) {
			t.Fatalf("fast writer stage metric missing %s:\n%s", want, text)
		}
	}
}

func TestProcessFastBatchAppendsTDengineBatchBeforeRedisAndAck(t *testing.T) {
	first := envelope.FrameEnvelope{Protocol: envelope.ProtocolJT808, VIN: "VIN001", EventID: "evt-4"}
	second := envelope.FrameEnvelope{Protocol: envelope.ProtocolJT808, VIN: "VIN002", EventID: "evt-5"}
	firstPayload, err := first.MarshalJSONBytes()
	if err != nil {
		t.Fatal(err)
	}
	secondPayload, err := second.MarshalJSONBytes()
	if err != nil {
		t.Fatal(err)
	}
	appender := &recordingFastAppender{}
	updater := &recordingFastUpdater{}
	ackCount := 0
	msgs := []*fastMessage{
		{subject: "vehicle.raw.go.jt808.v1", data: firstPayload, ack: func() error { ackCount++; return nil }},
		{subject: "vehicle.raw.go.jt808.v1", data: secondPayload, ack: func() error { ackCount++; return nil }},
	}

	if err := processFastBatch(context.Background(), nil, appender, updater, msgs); err != nil {
		t.Fatalf("processFastBatch() error = %v", err)
	}
	if appender.count != 0 {
		t.Fatalf("AppendAll count = %d, want 0", appender.count)
	}
	if appender.batchCount != 1 || appender.batchRows != 2 {
		t.Fatalf("AppendAllBatch count=%d rows=%d, want count=1 rows=2", appender.batchCount, appender.batchRows)
	}
	if updater.count != 2 || ackCount != 2 {
		t.Fatalf("updates=%d acks=%d, want 2/2", updater.count, ackCount)
	}
}

func TestProcessFastBatchExposesPendingMetricsDuringAppend(t *testing.T) {
	first := envelope.FrameEnvelope{Protocol: envelope.ProtocolJT808, VIN: "VIN001", EventID: "evt-6"}
	second := envelope.FrameEnvelope{Protocol: envelope.ProtocolJT808, VIN: "VIN002", EventID: "evt-7"}
	firstPayload, err := first.MarshalJSONBytes()
	if err != nil {
		t.Fatal(err)
	}
	secondPayload, err := second.MarshalJSONBytes()
	if err != nil {
		t.Fatal(err)
	}
	registry := metrics.NewRegistry()
	appenderStarted := make(chan struct{})
	appenderRelease := make(chan struct{})
	appender := &recordingFastAppender{
		onAppendBatch: func() {
			close(appenderStarted)
			<-appenderRelease
		},
	}
	msgs := []*fastMessage{
		{subject: "vehicle.raw.go.jt808.v1", data: firstPayload, ack: func() error { return nil }},
		{subject: "vehicle.raw.go.jt808.v1", data: secondPayload, ack: func() error { return nil }},
	}
	var wg sync.WaitGroup
	wg.Add(1)
	var processErr error
	go func() {
		defer wg.Done()
		processErr = processFastBatch(context.Background(), registry, appender, &recordingFastUpdater{}, msgs)
	}()
	<-appenderStarted

	text := registry.Render()
	for _, want := range []string{
		`vehicle_fast_writer_batch_pending_messages 2`,
		`vehicle_fast_writer_batch_pending_envelopes 2`,
	} {
		if !strings.Contains(text, want) {
			t.Fatalf("pending metric missing %s while appending:\n%s", want, text)
		}
	}

	close(appenderRelease)
	wg.Wait()
	if processErr != nil {
		t.Fatalf("processFastBatch() error = %v", processErr)
	}
	text = registry.Render()
	for _, want := range []string{
		`vehicle_fast_writer_batch_pending_messages 0`,
		`vehicle_fast_writer_batch_pending_envelopes 0`,
	} {
		if !strings.Contains(text, want) {
			t.Fatalf("pending metric should reset after append, missing %s:\n%s", want, text)
		}
	}
}

func TestRecordFastNATSConsumerInfoMetrics(t *testing.T) {
	registry := metrics.NewRegistry()
	cfg := config{NATSStream: "VEHICLE_INGEST", NATSDurable: "vehicle-fast-writer"}

	recordFastNATSConsumerInfoMetrics(registry, cfg, &nats.ConsumerInfo{
		NumPending:    23,
		NumAckPending: 5,
		NumWaiting:    3,
	})

	text := registry.Render()
	for _, want := range []string{
		`vehicle_fast_writer_nats_consumer_pending{consumer="vehicle-fast-writer",stream="VEHICLE_INGEST"} 23`,
		`vehicle_fast_writer_nats_consumer_ack_pending{consumer="vehicle-fast-writer",stream="VEHICLE_INGEST"} 5`,
		`vehicle_fast_writer_nats_consumer_waiting{consumer="vehicle-fast-writer",stream="VEHICLE_INGEST"} 3`,
	} {
		if !strings.Contains(text, want) {
			t.Fatalf("metrics missing %s:\n%s", want, text)
		}
	}
}

type recordingFastAppender struct {
	count         int
	batchCount    int
	batchRows     int
	err           error
	onAppendBatch func()
}

func (a *recordingFastAppender) AppendAll(context.Context, envelope.FrameEnvelope) error {
	a.count++
	return a.err
}

func (a *recordingFastAppender) AppendAllBatch(_ context.Context, envs []envelope.FrameEnvelope) error {
	a.batchCount++
	a.batchRows += len(envs)
	if a.onAppendBatch != nil {
		a.onAppendBatch()
	}
	return a.err
}

type recordingFastUpdater struct {
	count int
	err   error
}

func (u *recordingFastUpdater) FastUpdate(context.Context, envelope.FrameEnvelope) error {
	u.count++
	return u.err
}