handler-base/handler/handler_test.go

package handler

import (
	"context"
	"testing"

	"github.com/nats-io/nats.go"
	"github.com/vmihailenco/msgpack/v5"

	"git.daviestechlabs.io/daviestechlabs/handler-base/config"
)

// ────────────────────────────────────────────────────────────────────────────
// Handler construction tests
// ────────────────────────────────────────────────────────────────────────────

func TestNewHandler(t *testing.T) {
	cfg := config.Load()
	cfg.ServiceName = "test-handler"
	cfg.NATSQueueGroup = "test-group"

	h := New("ai.test.subject", cfg)
	if h.Subject != "ai.test.subject" {
		t.Errorf("Subject = %q", h.Subject)
	}
	if h.QueueGroup != "test-group" {
		t.Errorf("QueueGroup = %q", h.QueueGroup)
	}
	if h.Settings.ServiceName != "test-handler" {
		t.Errorf("ServiceName = %q", h.Settings.ServiceName)
	}
}

func TestNewHandlerNilSettings(t *testing.T) {
	h := New("ai.test", nil)
	if h.Settings == nil {
		t.Fatal("Settings should be loaded automatically")
	}
	if h.Settings.ServiceName != "handler" {
		t.Errorf("ServiceName = %q, want default", h.Settings.ServiceName)
	}
}

func TestCallbackRegistration(t *testing.T) {
	cfg := config.Load()
	h := New("ai.test", cfg)

	setupCalled := false
	h.OnSetup(func(ctx context.Context) error {
		setupCalled = true
		return nil
	})

	teardownCalled := false
	h.OnTeardown(func(ctx context.Context) error {
		teardownCalled = true
		return nil
	})

	h.OnMessage(func(ctx context.Context, msg *nats.Msg, data map[string]any) (map[string]any, error) {
		return nil, nil
	})

	if h.onSetup == nil || h.onTeardown == nil || h.onMessage == nil {
		t.Error("callbacks should not be nil after registration")
	}

	// Verify setup/teardown work when called directly.
	h.onSetup(context.Background())
	h.onTeardown(context.Background())
	if !setupCalled || !teardownCalled {
		t.Error("callbacks should have been invoked")
	}
}

// ────────────────────────────────────────────────────────────────────────────
// wrapHandler dispatch tests (unit test the message decode + dispatch logic)
// ────────────────────────────────────────────────────────────────────────────

func TestWrapHandler_ValidMessage(t *testing.T) {
	cfg := config.Load()
	h := New("ai.test", cfg)

	var receivedData map[string]any
	h.OnMessage(func(ctx context.Context, msg *nats.Msg, data map[string]any) (map[string]any, error) {
		receivedData = data
		return map[string]any{"status": "ok"}, nil
	})

	// Encode a message the same way services would.
	payload := map[string]any{
		"request_id": "test-001",
		"message":    "hello",
		"premium":    true,
	}
	encoded, err := msgpack.Marshal(payload)
	if err != nil {
		t.Fatal(err)
	}

	// Call wrapHandler directly without NATS.
	handler := h.wrapHandler(context.Background())
	handler(&nats.Msg{
		Subject: "ai.test.user.42.message",
		Data:    encoded,
	})

	if receivedData == nil {
		t.Fatal("handler was not called")
	}
	if receivedData["request_id"] != "test-001" {
		t.Errorf("request_id = %v", receivedData["request_id"])
	}
	if receivedData["premium"] != true {
		t.Errorf("premium = %v", receivedData["premium"])
	}
}

func TestWrapHandler_InvalidMsgpack(t *testing.T) {
	cfg := config.Load()
	h := New("ai.test", cfg)

	handlerCalled := false
	h.OnMessage(func(ctx context.Context, msg *nats.Msg, data map[string]any) (map[string]any, error) {
		handlerCalled = true
		return nil, nil
	})

	handler := h.wrapHandler(context.Background())
	handler(&nats.Msg{
		Subject: "ai.test",
		Data:    []byte{0xFF, 0xFE, 0xFD}, // invalid msgpack
	})

	if handlerCalled {
		t.Error("handler should not be called for invalid msgpack")
	}
}

func TestWrapHandler_HandlerError(t *testing.T) {
	cfg := config.Load()
	h := New("ai.test", cfg)

	h.OnMessage(func(ctx context.Context, msg *nats.Msg, data map[string]any) (map[string]any, error) {
		return nil, context.DeadlineExceeded
	})

	encoded, _ := msgpack.Marshal(map[string]any{"key": "val"})
	handler := h.wrapHandler(context.Background())

	// Should not panic even when handler returns error.
	handler(&nats.Msg{
		Subject: "ai.test",
		Data:    encoded,
	})
}

func TestWrapHandler_NilResponse(t *testing.T) {
	cfg := config.Load()
	h := New("ai.test", cfg)

	h.OnMessage(func(ctx context.Context, msg *nats.Msg, data map[string]any) (map[string]any, error) {
		return nil, nil // fire-and-forget style
	})

	encoded, _ := msgpack.Marshal(map[string]any{"x": 1})
	handler := h.wrapHandler(context.Background())

	// Should not panic with nil response and no reply subject.
	handler(&nats.Msg{
		Subject: "ai.test",
		Data:    encoded,
	})
}

// ────────────────────────────────────────────────────────────────────────────
// Benchmark: message decode + dispatch overhead
// ────────────────────────────────────────────────────────────────────────────

func BenchmarkWrapHandler(b *testing.B) {
	cfg := config.Load()
	h := New("ai.test", cfg)
	h.OnMessage(func(ctx context.Context, msg *nats.Msg, data map[string]any) (map[string]any, error) {
		return map[string]any{"ok": true}, nil
	})

	payload := map[string]any{
		"request_id": "bench-001",
		"message":    "What is the capital of France?",
		"premium":    true,
		"top_k":      10,
	}
	encoded, _ := msgpack.Marshal(payload)
	handler := h.wrapHandler(context.Background())
	msg := &nats.Msg{Subject: "ai.test", Data: encoded}

	b.ResetTimer()
	for b.Loop() {
		handler(msg)
	}
}