ExploreDNS/web/api/handler_internal_test.go

package api

import (
	"context"
	"net/http/httptest"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"
)

// TestStartTraversalJobCap verifies that new traversals are rejected with
// 429 once maxRunning jobs are in flight. The store is pre-filled with
// running jobs so no real traversal is spawned.
func TestStartTraversalJobCap(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	h := newHandler(ctx)
	h.maxRunning = 2

	for i := 0; i < 2; i++ {
		h.st.set(&TraversalJob{ID: strconv.Itoa(i), Status: statusRunning, StartedAt: time.Now()})
	}

	req := httptest.NewRequest("POST", "/api/traverse",
		strings.NewReader(`{"domain":"example.com"}`))
	rec := httptest.NewRecorder()
	h.mux.ServeHTTP(rec, req)

	if rec.Code != 429 {
		t.Fatalf("want 429 with %d running jobs, got %d: %s", h.st.runningCount(), rec.Code, rec.Body.String())
	}

	// A finished job frees a slot: runningCount must drop below the cap.
	done := time.Now()
	if j, ok := h.st.get("0"); ok {
		j.mu.Lock()
		j.Status = statusComplete
		j.DoneAt = &done
		j.mu.Unlock()
	}
	if got := h.st.runningCount(); got != 1 {
		t.Fatalf("runningCount after completion = %d, want 1", got)
	}
}

// TestJobTimeoutReachesTerminalState verifies that a traversal launched
// with an already-expired deadline still drives the job to a terminal
// error state (the TTL cleanup only ever purges finished jobs).
func TestJobTimeoutReachesTerminalState(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), time.Nanosecond)
	defer cancel()
	<-ctx.Done() // deadline already exceeded

	h := &Handler{st: newStore(), jobTimeout: time.Nanosecond}
	job := &TraversalJob{ID: "t", Status: statusRunning, StartedAt: time.Now(), cancel: cancel}
	h.st.set(job)

	doneCh := make(chan struct{})
	go func() {
		h.runTraversal(ctx, job, "example.com", 1, false)
		close(doneCh)
	}()
	select {
	case <-doneCh:
	case <-time.After(30 * time.Second):
		t.Fatal("runTraversal did not return with an expired context")
	}

	job.mu.RLock()
	defer job.mu.RUnlock()
	if job.DoneAt == nil {
		t.Fatal("job never reached a terminal state")
	}
	if job.Status != statusError || !strings.Contains(job.Error, "timed out") {
		t.Fatalf("want error status mentioning timeout, got status=%q error=%q", job.Status, job.Error)
	}
}

// TestSubscribeSnapshot_NoDuplicates regresses the subscribe/snapshot race:
// subscribers arriving while events are being published must never see the
// same event twice (once from the snapshot replay and once from the channel).
// Events are numbered, so any duplicate breaks strict monotonicity. Slow
// subscribers may legitimately drop events (publishLocked is non-blocking),
// so gaps are not an error.
func TestSubscribeSnapshot_NoDuplicates(t *testing.T) {
	job := &TraversalJob{Status: statusRunning}

	const total = 2000
	const subscribers = 8

	var pub sync.WaitGroup
	pub.Add(1)
	go func() {
		defer pub.Done()
		for i := 0; i < total; i++ {
			ev := ProgressEvent{RefID: strconv.Itoa(i)}
			job.mu.Lock()
			job.Progress = append(job.Progress, ev)
			job.publishLocked(ev)
			job.mu.Unlock()
		}
		job.mu.Lock()
		job.Status = statusComplete
		job.mu.Unlock()
		job.closeSubscribers()
	}()

	var subs sync.WaitGroup
	errs := make(chan string, subscribers)
	for s := 0; s < subscribers; s++ {
		subs.Add(1)
		go func() {
			defer subs.Done()
			sub, past, done := job.subscribeSnapshot()
			defer job.unsubscribe(sub)

			last := -1
			check := func(refid string) {
				n, err := strconv.Atoi(refid)
				if err != nil {
					errs <- "bad refid " + refid
					return
				}
				if n <= last {
					errs <- "event " + refid + " out of order or duplicated after " + strconv.Itoa(last)
					return
				}
				last = n
			}
			for _, ev := range past {
				check(ev.RefID)
			}
			if !done {
				for ev := range sub {
					check(ev.RefID)
				}
			}
		}()
	}

	pub.Wait()
	subs.Wait()
	close(errs)
	for msg := range errs {
		t.Error(msg)
	}
}