go-otel/internal/collector/collector_test.go

package collector

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	"git.tornberg.me/go-otel/internal/storage"
	"git.tornberg.me/go-otel/internal/websocket"
	"go.opentelemetry.io/collector/pdata/pcommon"
	"go.opentelemetry.io/collector/pdata/plog"
	"go.opentelemetry.io/collector/pdata/ptrace"
)

type noopBroadcaster struct{}

func (noopBroadcaster) Broadcast(entry *websocket.Entry) {}

// mockStore is used by tests only; addToHistory appends to the store
// asynchronously, so access must be synchronized even in single-threaded
// test flows that poll for the async write to land.
type mockStore struct {
	mu      sync.Mutex
	entries []websocket.HistoryEntry
}

func (m *mockStore) Append(ctx context.Context, entry websocket.HistoryEntry) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.entries = append(m.entries, entry)
	return nil
}

func (m *mockStore) len() int {
	m.mu.Lock()
	defer m.mu.Unlock()
	return len(m.entries)
}

func (m *mockStore) GetHistory(ctx context.Context, limit int) ([]websocket.HistoryEntry, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.entries, nil
}

func (m *mockStore) Search(ctx context.Context, opts storage.SearchOptions, limit int) ([]websocket.HistoryEntry, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var filtered []websocket.HistoryEntry
	for _, e := range m.entries {
		if opts.Type != "" && e.Type != opts.Type {
			continue
		}
		if opts.Query != "" && !strings.Contains(string(e.Data), opts.Query) {
			continue
		}
		filtered = append(filtered, e)
	}
	return filtered, nil
}

func TestCollectorWithStore(t *testing.T) {
	store := &mockStore{}
	coll := New(nil, 10, store)

	coll.addToHistory(websocket.DataTypeLogs, json.RawMessage(`{"msg": "test"}`))

	// Give it a moment since it's asynchronous
	timeout := time.After(100 * time.Millisecond)
	for {
		select {
		case <-timeout:
			t.Fatal("timed out waiting for entry to be appended to store")
		default:
			if store.len() > 0 {
				goto OK
			}
			time.Sleep(10 * time.Millisecond)
		}
	}

OK:
	entries, _ := store.GetHistory(context.Background(), 0)
	if entries[0].Type != "logs" {
		t.Errorf("expected type 'logs', got %s", entries[0].Type)
	}
}

func TestCollectorHistoryLimit(t *testing.T) {
	coll := New(nil, 10, nil)

	for i := 0; i < 10; i++ {
		coll.addToHistory(websocket.DataTypeLogs, json.RawMessage(`{"msg": "test"}`))
	}

	history := coll.GetHistory(10)
	if len(history) != 10 {
		t.Errorf("expected 10 entries, got %d", len(history))
	}

	history = coll.GetHistory(5)
	if len(history) != 5 {
		t.Errorf("expected 5 entries, got %d", len(history))
	}

	// Should default to historySize (10)
	history = coll.GetHistory(0)
	if len(history) != 10 {
		t.Errorf("expected default limit to be historySize (10), got %d", len(history))
	}

	// Should not crash on large limit, just return what it has
	history = coll.GetHistory(100)
	if len(history) != 10 {
		t.Errorf("expected 10 entries even with limit 100, got %d", len(history))
	}
}

func TestHandleHistoryByTypeUsesStore(t *testing.T) {
	store := &mockStore{
		entries: []websocket.HistoryEntry{
			{Type: "traces", Timestamp: 1, Data: json.RawMessage(`{"id":1}`)},
			{Type: "traces", Timestamp: 2, Data: json.RawMessage(`{"id":2}`)},
			{Type: "logs", Timestamp: 3, Data: json.RawMessage(`{"id":3}`)},
		},
	}
	// Ring buffer is much smaller than the store's content, so a correct
	// implementation must route through the store rather than the ring buffer.
	coll := New(nil, 1, store)

	req := httptest.NewRequest(http.MethodGet, "/api/history/traces", nil)
	rec := httptest.NewRecorder()
	coll.HandleHistoryByType("traces")(rec, req)

	var results []websocket.HistoryEntry
	if err := json.Unmarshal(rec.Body.Bytes(), &results); err != nil {
		t.Fatalf("failed to decode response: %v", err)
	}
	if len(results) != 2 {
		t.Fatalf("expected 2 trace entries from the store, got %d", len(results))
	}
}

func TestHandleTraceByID(t *testing.T) {
	store := &mockStore{
		entries: []websocket.HistoryEntry{
			{Type: "traces", Timestamp: 1, Data: json.RawMessage(`{"traceId":"abc123"}`)},
			{Type: "logs", Timestamp: 2, Data: json.RawMessage(`{"traceId":"abc123","msg":"related"}`)},
			{Type: "traces", Timestamp: 3, Data: json.RawMessage(`{"traceId":"other"}`)},
		},
	}
	coll := New(nil, 10, store)

	req := httptest.NewRequest(http.MethodGet, "/api/traces/abc123", nil)
	req.SetPathValue("traceId", "abc123")
	rec := httptest.NewRecorder()
	coll.HandleTraceByID(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d: %s", rec.Code, rec.Body.String())
	}

	var results []websocket.HistoryEntry
	if err := json.Unmarshal(rec.Body.Bytes(), &results); err != nil {
		t.Fatalf("failed to decode response: %v", err)
	}
	if len(results) != 2 {
		t.Fatalf("expected 2 entries referencing trace abc123, got %d", len(results))
	}
}

func TestHandleWaitRateLimited(t *testing.T) {
	coll := New(nil, 10, nil)

	var lastCode int
	for i := 0; i < 10; i++ {
		req := httptest.NewRequest(http.MethodGet, "/api/history/wait?q=nope&timeout=0", nil)
		req.RemoteAddr = "203.0.113.1:12345"
		rec := httptest.NewRecorder()
		coll.HandleWait(rec, req)
		lastCode = rec.Code
		if lastCode == http.StatusTooManyRequests {
			break
		}
	}
	if lastCode != http.StatusTooManyRequests {
		t.Fatalf("expected rate limiting to eventually return 429, last code was %d", lastCode)
	}
}

func TestConsumeTracesRecordsIssue(t *testing.T) {
	coll := New(noopBroadcaster{}, 10, nil)

	td := ptrace.NewTraces()
	rs := td.ResourceSpans().AppendEmpty()
	rs.Resource().Attributes().PutStr("service.name", "checkout")
	span := rs.ScopeSpans().AppendEmpty().Spans().AppendEmpty()
	span.SetName("charge-card")
	span.Status().SetCode(ptrace.StatusCodeError)
	span.Status().SetMessage("card declined")
	span.SetTraceID(pcommon.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16})

	if err := coll.ConsumeTraces(context.Background(), td); err != nil {
		t.Fatalf("ConsumeTraces failed: %v", err)
	}

	list := coll.ListIssues(10)
	if len(list) != 1 {
		t.Fatalf("expected 1 issue, got %d", len(list))
	}
	if list[0].Service != "checkout" {
		t.Errorf("expected service 'checkout', got %s", list[0].Service)
	}
	if list[0].Kind != "span_error" {
		t.Errorf("expected kind 'span_error', got %s", list[0].Kind)
	}
	if !strings.Contains(list[0].Title, "card declined") {
		t.Errorf("expected title to contain status message, got %q", list[0].Title)
	}
	if list[0].Samples[0].TraceID == "" {
		t.Errorf("expected sample to carry the trace id")
	}

	// A non-error span must not create an issue.
	okSpan := rs.ScopeSpans().At(0).Spans().AppendEmpty()
	okSpan.SetName("healthy-op")
	okSpan.Status().SetCode(ptrace.StatusCodeOk)
	if err := coll.ConsumeTraces(context.Background(), td); err != nil {
		t.Fatalf("ConsumeTraces failed: %v", err)
	}

	// Re-sending the same error span dedupes into the existing issue.
	list = coll.ListIssues(10)
	if len(list) != 1 || list[0].Count != 2 {
		t.Fatalf("expected deduped issue with count 2, got %+v", list)
	}
}

func TestConsumeLogsRecordsIssue(t *testing.T) {
	coll := New(noopBroadcaster{}, 10, nil)

	ld := plog.NewLogs()
	rl := ld.ResourceLogs().AppendEmpty()
	rl.Resource().Attributes().PutStr("service.name", "api")
	record := rl.ScopeLogs().AppendEmpty().LogRecords().AppendEmpty()
	record.SetSeverityNumber(plog.SeverityNumberError)
	record.Body().SetStr("panic: nil pointer dereference")

	// An info-level record must not create an issue.
	infoRecord := rl.ScopeLogs().At(0).LogRecords().AppendEmpty()
	infoRecord.SetSeverityNumber(plog.SeverityNumberInfo)
	infoRecord.Body().SetStr("request handled")

	if err := coll.ConsumeLogs(context.Background(), ld); err != nil {
		t.Fatalf("ConsumeLogs failed: %v", err)
	}

	list := coll.ListIssues(10)
	if len(list) != 1 {
		t.Fatalf("expected 1 issue, got %d", len(list))
	}
	if list[0].Kind != "log_error" {
		t.Errorf("expected kind 'log_error', got %s", list[0].Kind)
	}
	if list[0].Title != "panic: nil pointer dereference" {
		t.Errorf("unexpected title: %q", list[0].Title)
	}
}

func TestHandleIssuesAndByID(t *testing.T) {
	coll := New(noopBroadcaster{}, 10, nil)

	ld := plog.NewLogs()
	rl := ld.ResourceLogs().AppendEmpty()
	rl.Resource().Attributes().PutStr("service.name", "api")
	record := rl.ScopeLogs().AppendEmpty().LogRecords().AppendEmpty()
	record.SetSeverityNumber(plog.SeverityNumberError)
	record.Body().SetStr("boom")

	if err := coll.ConsumeLogs(context.Background(), ld); err != nil {
		t.Fatalf("ConsumeLogs failed: %v", err)
	}

	req := httptest.NewRequest(http.MethodGet, "/api/issues", nil)
	rec := httptest.NewRecorder()
	coll.HandleIssues(rec, req)

	var list []map[string]any
	if err := json.Unmarshal(rec.Body.Bytes(), &list); err != nil {
		t.Fatalf("failed to decode issues response: %v", err)
	}
	if len(list) != 1 {
		t.Fatalf("expected 1 issue in response, got %d", len(list))
	}
	id, _ := list[0]["id"].(string)
	if id == "" {
		t.Fatalf("expected issue to have an id, got %+v", list[0])
	}

	req = httptest.NewRequest(http.MethodGet, "/api/issues/"+id, nil)
	req.SetPathValue("id", id)
	rec = httptest.NewRecorder()
	coll.HandleIssueByID(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("expected 200 for known issue id, got %d", rec.Code)
	}

	req = httptest.NewRequest(http.MethodGet, "/api/issues/missing", nil)
	req.SetPathValue("id", "missing")
	rec = httptest.NewRecorder()
	coll.HandleIssueByID(rec, req)
	if rec.Code != http.StatusNotFound {
		t.Fatalf("expected 404 for unknown issue id, got %d", rec.Code)
	}
}

func TestHandleLogsIngestKey(t *testing.T) {
	coll := New(noopBroadcaster{}, 10, nil)
	coll.SetIngestKey("secret")

	body := []byte(`{"resourceLogs":[]}`)

	cases := []struct {
		name       string
		url        string
		header     string
		wantStatus int
	}{
		{"no key", "/v1/logs", "", http.StatusUnauthorized},
		{"wrong key", "/v1/logs?key=wrong", "", http.StatusUnauthorized},
		{"correct key in query", "/v1/logs?key=secret", "", http.StatusOK},
		{"correct key in header", "/v1/logs", "secret", http.StatusOK},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodPost, tc.url, bytes.NewReader(body))
			req.Header.Set("Content-Type", "application/json")
			if tc.header != "" {
				req.Header.Set("X-Otel-Ingest-Key", tc.header)
			}
			// Shared IP across subtests; the ingest limiter's burst of 50 is
			// well above this test's request count.
			req.RemoteAddr = "198.51.100.1:1234"
			rec := httptest.NewRecorder()
			coll.HandleLogs(rec, req)
			if rec.Code != tc.wantStatus {
				t.Errorf("expected status %d, got %d (body=%s)", tc.wantStatus, rec.Code, rec.Body.String())
			}
		})
	}
}

func TestHandleLogsIngestRateLimited(t *testing.T) {
	coll := New(noopBroadcaster{}, 10, nil)
	body := []byte(`{"resourceLogs":[]}`)

	var lastCode int
	for i := 0; i < 80; i++ {
		req := httptest.NewRequest(http.MethodPost, "/v1/logs", bytes.NewReader(body))
		req.Header.Set("Content-Type", "application/json")
		req.RemoteAddr = "198.51.100.2:1234"
		rec := httptest.NewRecorder()
		coll.HandleLogs(rec, req)
		lastCode = rec.Code
		if lastCode == http.StatusTooManyRequests {
			break
		}
	}
	if lastCode != http.StatusTooManyRequests {
		t.Fatalf("expected ingest rate limiting to eventually return 429, last code was %d", lastCode)
	}
}