fetch_ml/tests/unit/logging/logging_test.go

package tests

import (
	"context"
	"io"
	"log/slog"
	"os"
	"os/exec"
	"strings"
	"testing"

	"github.com/jfraeys/fetch_ml/internal/logging"
)

type recordingHandler struct {
	base []slog.Attr
	last []slog.Attr
}

func TestLoggerFatalExits(t *testing.T) {
	if os.Getenv("LOG_FATAL_TEST") == "1" {
		logger := logging.NewLogger(slog.LevelInfo, false)
		logger.Fatal("fatal message")
		return
	}

	cmd := exec.Command(os.Args[0], "-test.run", t.Name())
	cmd.Env = append(os.Environ(), "LOG_FATAL_TEST=1")
	if err := cmd.Run(); err == nil {
		t.Fatalf("expected Fatal to exit with non-nil error")
	}
}

func TestNewLoggerHonorsJSONFormatEnv(t *testing.T) {
	t.Setenv("LOG_FORMAT", "json")
	origStderr := os.Stderr
	r, w, err := os.Pipe()
	if err != nil {
		t.Fatalf("failed to create pipe: %v", err)
	}
	os.Stderr = w
	defer func() {
		w.Close()
		r.Close()
		os.Stderr = origStderr
	}()

	logger := logging.NewLogger(slog.LevelInfo, false)
	logger.Info("hello", "key", "value")
	w.Close()
	data, readErr := io.ReadAll(r)
	if readErr != nil {
		t.Fatalf("failed to read logger output: %v", readErr)
	}

	output := string(data)
	if !strings.Contains(output, "\"msg\":\"hello\"") || !strings.Contains(output, "\"key\":\"value\"") {
		t.Fatalf("expected json output, got %s", output)
	}
}

func (h *recordingHandler) Enabled(_ context.Context, _ slog.Level) bool {
	return true
}

func (h *recordingHandler) Handle(_ context.Context, r slog.Record) error {
	// Reset last and include base attributes first
	h.last = append([]slog.Attr{}, h.base...)
	r.Attrs(func(a slog.Attr) bool {
		h.last = append(h.last, a)
		return true
	})
	return nil
}

func (h *recordingHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	newBase := append([]slog.Attr{}, h.base...)
	newBase = append(newBase, attrs...)
	return &recordingHandler{base: newBase}
}

func (h *recordingHandler) WithGroup(_ string) slog.Handler {
	return h
}

func attrsToMap(attrs []slog.Attr) map[string]any {
	out := make(map[string]any, len(attrs))
	for _, attr := range attrs {
		out[attr.Key] = attr.Value.Any()
	}
	return out
}

func TestEnsureTraceAddsIDs(t *testing.T) {
	ctx := context.Background()
	ctx = logging.EnsureTrace(ctx)

	if ctx.Value(logging.CtxTraceID) == nil {
		t.Fatalf("expected trace id to be injected")
	}
	if ctx.Value(logging.CtxSpanID) == nil {
		t.Fatalf("expected span id to be injected")
	}

	existingTrace := ctx.Value(logging.CtxTraceID)
	ctx = logging.EnsureTrace(ctx)
	if ctx.Value(logging.CtxTraceID) != existingTrace {
		t.Fatalf("EnsureTrace should not overwrite existing trace id")
	}
}

func TestLoggerWithContextIncludesValues(t *testing.T) {
	handler := &recordingHandler{}
	base := slog.New(handler)
	logger := &logging.Logger{Logger: base}

	ctx := context.Background()
	ctx = context.WithValue(ctx, logging.CtxTraceID, "trace-123")
	ctx = context.WithValue(ctx, logging.CtxSpanID, "span-456")
	ctx = logging.CtxWithWorker(ctx, "worker-1")
	ctx = logging.CtxWithJob(ctx, "job-a")
	ctx = logging.CtxWithTask(ctx, "task-b")

	child := logger.WithContext(ctx)
	child.Info("hello")

	rec, ok := child.Handler().(*recordingHandler)
	if !ok {
		t.Fatalf("expected recordingHandler, got %T", child.Handler())
	}

	fields := attrsToMap(rec.last)
	expected := map[string]string{
		"trace_id":  "trace-123",
		"span_id":   "span-456",
		"worker_id": "worker-1",
		"job_name":  "job-a",
		"task_id":   "task-b",
	}

	for key, want := range expected {
		got, ok := fields[key]
		if !ok {
			t.Fatalf("expected attribute %s to be present", key)
		}
		if got != want {
			t.Fatalf("attribute %s mismatch: want %s got %v", key, want, got)
		}
	}
}

func TestColorTextHandlerAddsColorAttr(t *testing.T) {
	tmp, err := os.CreateTemp("", "log-output")
	if err != nil {
		t.Fatalf("failed to create temp file: %v", err)
	}
	t.Cleanup(func() {
		tmp.Close()
		os.Remove(tmp.Name())
	})

	handler := logging.NewColorTextHandler(tmp, &slog.HandlerOptions{Level: slog.LevelInfo})
	logger := slog.New(handler)

	logger.Info("color test")

	if err := tmp.Sync(); err != nil {
		t.Fatalf("failed to sync temp file: %v", err)
	}

	data, err := os.ReadFile(tmp.Name())
	if err != nil {
		t.Fatalf("failed to read temp file: %v", err)
	}

	output := string(data)
	if !strings.Contains(output, "lvl_color=\"\x1b[32mINF\x1b[0m\"") &&
		!strings.Contains(output, "lvl_color=\"\\x1b[32mINF\\x1b[0m\"") {
		t.Fatalf("expected info level color attribute, got: %s", output)
	}
}