iop/apps/node/internal/node/health_probe_test.go

package node

import (
	"context"
	"errors"
	"testing"
	"time"

	runtime "iop/packages/go/execution"
)

// recordingProbe captures the context the coordinator passed to the probe hook
// so tests can assert it is live, independent, and exactly bounded.
type recordingProbe struct {
	ctx     context.Context
	result  runtime.ProviderProbeResult
	err     error
	calls   int
	probeFn func(ctx context.Context, target string) (runtime.ProviderProbeResult, error)
}

func (r *recordingProbe) probe(ctx context.Context, target string) (runtime.ProviderProbeResult, error) {
	r.calls++
	r.ctx = ctx
	if r.probeFn != nil {
		return r.probeFn(ctx, target)
	}
	return r.result, r.err
}

func TestProbeHealthAvailableYieldsRequestStalled(t *testing.T) {
	rec := &recordingProbe{result: runtime.ProviderProbeResult{
		AdapterName: "vllm", InstanceKey: "vllm-gpu", Target: "m-a",
		Status: runtime.ProviderStatusAvailable,
	}}
	ev := ProbeHealth("vllm", "vllm-gpu", "m-a", rec.probe)
	if ev.Health != runtime.RequestStalled {
		t.Fatalf("Health: got %q, want %q", ev.Health, runtime.RequestStalled)
	}
	if ev.Status != runtime.ProviderStatusAvailable {
		t.Errorf("Status: got %q, want available", ev.Status)
	}
	if rec.calls != 1 {
		t.Errorf("probe called %d times, want 1", rec.calls)
	}
}

func TestProbeHealthUnavailableYieldsProviderUnhealthy(t *testing.T) {
	rec := &recordingProbe{result: runtime.ProviderProbeResult{
		AdapterName: "ollama", Target: "m-b",
		Status: runtime.ProviderStatusUnavailable,
	}}
	ev := ProbeHealth("ollama", "", "m-b", rec.probe)
	if ev.Health != runtime.ProviderUnhealthy {
		t.Fatalf("Health: got %q, want %q", ev.Health, runtime.ProviderUnhealthy)
	}
	if ev.Status != runtime.ProviderStatusUnavailable {
		t.Errorf("Status: got %q, want unavailable", ev.Status)
	}
}

func TestProbeHealthTransportErrorYieldsHealthUnknown(t *testing.T) {
	boom := errors.New("connection refused")
	rec := &recordingProbe{
		result: runtime.ProviderProbeResult{AdapterName: "vllm", Target: "m-a"},
		err:    boom,
	}
	ev := ProbeHealth("vllm", "", "m-a", rec.probe)
	if ev.Health != runtime.HealthUnknown {
		t.Fatalf("Health: got %q, want %q", ev.Health, runtime.HealthUnknown)
	}
	if ev.Status != runtime.ProviderStatusUnknown {
		t.Errorf("Status: got %q, want unknown", ev.Status)
	}
	if ev.Detail != boom.Error() {
		t.Errorf("Detail: got %q, want %q", ev.Detail, boom.Error())
	}
}

func TestProbeHealthDeadlineExceededYieldsHealthUnknown(t *testing.T) {
	rec := &recordingProbe{
		result: runtime.ProviderProbeResult{AdapterName: "vllm", Target: "m-a"},
		err:    context.DeadlineExceeded,
	}
	ev := ProbeHealth("vllm", "", "m-a", rec.probe)
	if ev.Health != runtime.HealthUnknown {
		t.Fatalf("Health: got %q, want %q", ev.Health, runtime.HealthUnknown)
	}
	if ev.Detail != "probe timed out" {
		t.Errorf("Detail: got %q, want probe timed out", ev.Detail)
	}
}

func TestProbeHealthUnsupportedAdapterYieldsHealthUnknown(t *testing.T) {
	ev := ProbeHealth("worker", "", "m-a", nil)
	if ev.Health != runtime.HealthUnknown {
		t.Fatalf("Health: got %q, want %q", ev.Health, runtime.HealthUnknown)
	}
	if ev.Detail != "adapter does not support probing" {
		t.Errorf("Detail: got %q", ev.Detail)
	}
}

func TestProbeHealthIdentityMismatchYieldsHealthUnknown(t *testing.T) {
	// The probe confirms a different adapter/target than the request required.
	rec := &recordingProbe{result: runtime.ProviderProbeResult{
		AdapterName: "ollama", Target: "m-a",
		Status: runtime.ProviderStatusAvailable,
	}}
	ev := ProbeHealth("vllm", "", "m-a", rec.probe)
	if ev.Health != runtime.HealthUnknown {
		t.Fatalf("Health: got %q, want %q", ev.Health, runtime.HealthUnknown)
	}
	if ev.Detail != "probe identity did not match request" {
		t.Errorf("Detail: got %q", ev.Detail)
	}
}

func TestProbeHealthPinnedInstanceMismatchYieldsHealthUnknown(t *testing.T) {
	rec := &recordingProbe{result: runtime.ProviderProbeResult{
		AdapterName: "vllm", InstanceKey: "vllm-gpu", Target: "m-a",
		Status: runtime.ProviderStatusAvailable,
	}}
	ev := ProbeHealth("vllm", "vllm-other", "m-a", rec.probe)
	if ev.Health != runtime.HealthUnknown {
		t.Fatalf("Health: got %q, want %q", ev.Health, runtime.HealthUnknown)
	}
}

// TestProbeHealthRechecksDeadlineWhenProbeIgnoresContext proves the coordinator
// re-checks its independent deadline after the probe returns nil error. The
// ceiling is lowered to the past so the rooted context is already expired; a
// probe that ignores that context and reports available must still be
// classified inconclusive. No scheduler sleep is used.
func TestProbeHealthRechecksDeadlineWhenProbeIgnoresContext(t *testing.T) {
	saved := healthProbeCeiling
	healthProbeCeiling = -1 * time.Millisecond
	defer func() { healthProbeCeiling = saved }()

	rec := &recordingProbe{result: runtime.ProviderProbeResult{
		AdapterName: "vllm", Target: "m-a",
		Status: runtime.ProviderStatusAvailable,
	}}
	ev := ProbeHealth("vllm", "", "m-a", rec.probe)
	if ev.Health != runtime.HealthUnknown {
		t.Fatalf("Health: got %q, want %q after ignored deadline", ev.Health, runtime.HealthUnknown)
	}
	if ev.Detail != "probe timed out" {
		t.Errorf("Detail: got %q, want probe timed out", ev.Detail)
	}
}

// TestProbeHealthReturnsWhenBlockedHookOutlivesContext proves the coordinator
// returns at its independent ceiling even when the prober ignores context
// cancellation and never returns. It exercises the unexported context-taking
// runProbe helper with a manually canceled context: the hook signals started,
// the test cancels the context, the coordinator must return fail-closed
// (health_unknown / probe timed out) while the hook is still blocked, and only
// then does the test release the hook so no goroutine leaks. No time.Sleep,
// wall-clock polling, live provider, or arbitrary provider metadata is used.
func TestProbeHealthReturnsWhenBlockedHookOutlivesContext(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	started := make(chan struct{})
	release := make(chan struct{})
	done := make(chan runtime.ProbeOutcome, 1)

	probe := func(_ context.Context, _ string) (runtime.ProviderProbeResult, error) {
		started <- struct{}{}
		<-release
		return runtime.ProviderProbeResult{}, nil
	}

	go func() {
		done <- runProbe(ctx, "m-a", probe, runtime.ProbeOutcome{
			ExpectedAdapter: "vllm",
			ExpectedTarget:  "m-a",
		})
	}()

	<-started

	// Cancel the manual context. The coordinator must return fail-closed while
	// the hook is still blocked on release.
	cancel()

	select {
	case got := <-done:
		if class := runtime.ClassifyProbeOutcome(got); class != runtime.LivenessTimeout {
			t.Fatalf("classification: got %q, want %q", class, runtime.LivenessTimeout)
		}
		ev := finalizeHealthProbe(got)
		if ev.Health != runtime.HealthUnknown {
			t.Fatalf("Health: got %q, want %q while hook still blocked", ev.Health, runtime.HealthUnknown)
		}
		if ev.Detail != "probe timed out" {
			t.Errorf("Detail: got %q, want probe timed out", ev.Detail)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("coordinator did not return within 2s after context cancel; hook held it past the ceiling")
	}

	// Release the blocked hook so the probe goroutine finishes and no goroutine
	// leaks past the test.
	close(release)
}

// TestProbeHealthReceivesIndependentBoundedContext proves the probe hook
// receives a live, independently rooted, exactly bounded context: it has its
// own deadline near the ceiling and is not derived from any canceled execution
// request (the coordinator takes no execution context by design). The context
// state is snapshotted inside the probe hook because ProbeHealth cancels its
// rooted context after returning.
func TestProbeHealthReceivesIndependentBoundedContext(t *testing.T) {
	var (
		observedAt       time.Time
		observedDeadline time.Time
		hasDeadline      bool
		observedErr      error
		observedPtr      interface{ Done() <-chan struct{} }
	)
	probe := func(ctx context.Context, target string) (runtime.ProviderProbeResult, error) {
		observedAt = time.Now()
		observedDeadline, hasDeadline = ctx.Deadline()
		observedErr = ctx.Err()
		observedPtr = ctx
		return runtime.ProviderProbeResult{
			AdapterName: "vllm", Target: "m-a",
			Status: runtime.ProviderStatusAvailable,
		}, nil
	}
	_ = ProbeHealth("vllm", "", "m-a", probe)

	if observedErr != nil {
		t.Fatalf("probe context not live: %v", observedErr)
	}
	if !hasDeadline {
		t.Fatal("probe context has no deadline")
	}
	if !observedDeadline.After(observedAt) {
		t.Fatalf("probe deadline %v is not in the future (now %v)", observedDeadline, observedAt)
	}
	if got := observedDeadline.Sub(observedAt); got > healthProbeCeiling {
		t.Fatalf("probe bound %v exceeds ceiling %v", got, healthProbeCeiling)
	}
	// The rooted context must not be tied to a caller-supplied context.
	cancelCtx, cancel := context.WithCancel(context.Background())
	cancel()
	if observedPtr == cancelCtx {
		t.Fatal("probe context must not be a caller-supplied context")
	}
}

// TestProbeHealthRootsFromBackground proves a cancelled caller-side context
// cannot cut the probe short: the coordinator takes no execution context by
// design, so the probe still observes a live, bounded context and a definitive
// result despite an unrelated canceled context existing in the caller.
func TestProbeHealthRootsFromBackground(t *testing.T) {
	saved := healthProbeCeiling
	healthProbeCeiling = 50 * time.Millisecond
	defer func() { healthProbeCeiling = saved }()

	// A separate canceled context exists in the caller; the coordinator must
	// not be derived from it.
	_, cancel := context.WithCancel(context.Background())
	cancel()

	var observedErr error
	probe := func(ctx context.Context, target string) (runtime.ProviderProbeResult, error) {
		observedErr = ctx.Err()
		return runtime.ProviderProbeResult{
			AdapterName: "vllm", Target: "m-a",
			Status: runtime.ProviderStatusAvailable,
		}, nil
	}
	ev := ProbeHealth("vllm", "", "m-a", probe)
	if ev.Health != runtime.RequestStalled {
		t.Fatalf("Health: got %q, want %q (caller cancellation must not affect probe)", ev.Health, runtime.RequestStalled)
	}
	if observedErr != nil {
		t.Fatalf("probe context was not live despite a canceled caller-side context: %v", observedErr)
	}
}

type stubProberProvider struct {
	probed bool
}

func (s *stubProberProvider) Name() string { return "stub" }
func (s *stubProberProvider) Capabilities(_ context.Context) (runtime.Capabilities, error) {
	return runtime.Capabilities{AdapterName: "stub"}, nil
}
func (s *stubProberProvider) Execute(_ context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
	return nil
}
func (s *stubProberProvider) ProbeProvider(_ context.Context, _ string) (runtime.ProviderProbeResult, error) {
	s.probed = true
	return runtime.ProviderProbeResult{AdapterName: "stub", Target: "m-a", Status: runtime.ProviderStatusAvailable}, nil
}

type stubPlainProvider struct{}

func (s *stubPlainProvider) Name() string { return "plain" }
func (s *stubPlainProvider) Capabilities(_ context.Context) (runtime.Capabilities, error) {
	return runtime.Capabilities{AdapterName: "plain"}, nil
}
func (s *stubPlainProvider) Execute(_ context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
	return nil
}

func TestResolveProbeFunc(t *testing.T) {
	t.Run("prober_adapter_returns_hook", func(t *testing.T) {
		stub := &stubProberProvider{}
		probe := ResolveProbeFunc(stub)
		if probe == nil {
			t.Fatal("expected non-nil probe hook for prober adapter")
		}
		res, err := probe(context.Background(), "m-a")
		if err != nil || res.Status != runtime.ProviderStatusAvailable {
			t.Fatalf("unexpected probe result: %+v err=%v", res, err)
		}
		if !stub.probed {
			t.Fatal("probe hook did not invoke ProviderProber.ProbeProvider")
		}
	})
	t.Run("plain_adapter_returns_nil", func(t *testing.T) {
		if ResolveProbeFunc(&stubPlainProvider{}) != nil {
			t.Fatal("expected nil probe hook for non-prober adapter")
		}
	})
}

func TestProbeHealthViaResolveProbeFuncEndToEnd(t *testing.T) {
	ev := ProbeHealth("stub", "", "m-a", ResolveProbeFunc(&stubProberProvider{}))
	if ev.Health != runtime.RequestStalled {
		t.Fatalf("Health: got %q, want %q", ev.Health, runtime.RequestStalled)
	}
}