pangolin/server/internal/scheduler/probe/prober_agent_test.go

package probe_test

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

	"github.com/alicebob/miniredis/v2"
	"github.com/redis/go-redis/v9"
	"github.com/wangjia/pangolin/server/internal/scheduler/probe"
)

// --------------------------------------------------------------------------
// Mock vendor server helpers
// --------------------------------------------------------------------------

// vendorBehavior controls what the mock Alibaba Cloud server returns.
type vendorBehavior int

const (
	behaviorSuccess     vendorBehavior = iota // normal TCP probe success
	behaviorFailedProbe                       // probe ran but target unreachable
	behaviorRateLimit                         // HTTP 429
	behaviorTimeout                           // server hangs until context cancelled
)

// newMockAliyunServer creates an httptest.Server that behaves according to b.
// It counts how many times CreateSiteMonitor and DescribeSiteMonitorData were called.
func newMockAliyunServer(t *testing.T, b vendorBehavior) (*httptest.Server, *atomic.Int32, *atomic.Int32) {
	t.Helper()
	var createCalls, pollCalls atomic.Int32

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := r.ParseForm(); err != nil {
			http.Error(w, "bad form", 400)
			return
		}
		action := r.FormValue("Action")

		switch b {
		case behaviorRateLimit:
			w.WriteHeader(http.StatusTooManyRequests)
			return

		case behaviorTimeout:
			// Block until the client closes or test ends.
			<-r.Context().Done()
			return
		}

		switch action {
		case "CreateSiteMonitor":
			createCalls.Add(1)
			if b == behaviorSuccess || b == behaviorFailedProbe {
				writeJSON(w, map[string]interface{}{
					"Code":      "200",
					"Message":   "OK",
					"RequestId": "mock-req-id",
					"Data":      map[string]string{"TaskId": "task-mock-001"},
				})
			}

		case "DescribeSiteMonitorData":
			pollCalls.Add(1)
			switch b {
			case behaviorSuccess:
				writeJSON(w, map[string]interface{}{
					"Code":    "200",
					"Message": "OK",
					"Data": map[string]interface{}{
						"List": []map[string]interface{}{
							{
								"Province":  "北京",
								"Average":   float64(42),
								"Value":     float64(100), // 100 = available
								"ErrorCode": "",
								"ErrorInfo": "",
							},
						},
					},
				})
			case behaviorFailedProbe:
				writeJSON(w, map[string]interface{}{
					"Code":    "200",
					"Message": "OK",
					"Data": map[string]interface{}{
						"List": []map[string]interface{}{
							{
								"Province":  "北京",
								"Average":   float64(0),
								"Value":     float64(0), // 0 = unavailable
								"ErrorCode": "CONNECT_TIMEOUT",
								"ErrorInfo": "connect timeout",
							},
						},
					},
				})
			}

		default:
			w.WriteHeader(http.StatusBadRequest)
		}
	}))

	t.Cleanup(srv.Close)
	return srv, &createCalls, &pollCalls
}

func writeJSON(w http.ResponseWriter, v interface{}) {
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(v)
}

// newTestAgent creates an AliyunSyntheticAgent wired to a mock server and miniredis.
func newTestAgent(t *testing.T, srv *httptest.Server, mr *miniredis.Miniredis) *probe.AliyunSyntheticAgent {
	t.Helper()
	rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
	st := probe.NewStore(rdb)

	cfg := probe.AliyunSyntheticAgentConfig{
		AccessKeyID:     "test-key-id",
		AccessKeySecret: "test-key-secret",
		PollInterval:    10 * time.Millisecond, // fast polling in tests
		PollTimeout:     2 * time.Second,
		ProbeCity:       "563",
	}

	var logBuf bytes.Buffer
	logger := slog.New(slog.NewJSONHandler(&logBuf, nil))

	agent := probe.NewAliyunSyntheticAgentWithOptions(cfg, st, logger, srv.URL+"/")
	return agent
}

// --------------------------------------------------------------------------
// Tests
// --------------------------------------------------------------------------

// TestProbeSuccess_NormalizesResults verifies that a successful vendor response
// is correctly mapped to VantageResults with 3rd- prefixed ISP names.
func TestProbeSuccess_NormalizesResults(t *testing.T) {
	mr := miniredis.RunT(t)
	srv, createCalls, _ := newMockAliyunServer(t, behaviorSuccess)
	agent := newTestAgent(t, srv, mr)

	target := probe.ProbeTarget{NodeID: "node-hk-01", Host: "hk.example.com", Port: 443}
	results, err := agent.Probe(context.Background(), target)
	if err != nil {
		t.Fatalf("Probe returned error: %v", err)
	}

	// Expect one result per ISP (3 total: Telecom, Unicom, Mobile).
	if len(results) != 3 {
		t.Fatalf("expected 3 VantageResults, got %d", len(results))
	}

	// All ISP names must carry the "3rd-" prefix.
	for _, vr := range results {
		if !strings.HasPrefix(vr.Vantage.ISP, "3rd-") {
			t.Errorf("ISP %q missing 3rd- prefix", vr.Vantage.ISP)
		}
		if vr.Vantage.Country != "CN" {
			t.Errorf("Country = %q, want CN", vr.Vantage.Country)
		}
	}

	// Check first result details (mock returns same data for all ISPs).
	vr := results[0]
	if !vr.L1.OK {
		t.Error("L1.OK should be true on success")
	}
	if vr.L1.RttMs != 42 {
		t.Errorf("L1.RttMs = %d, want 42", vr.L1.RttMs)
	}
	if vr.L2 == nil {
		t.Fatal("L2 should be non-nil for port 443")
	}
	if !vr.L2.OK {
		t.Errorf("L2.OK should be true; Err=%q", vr.L2.Err)
	}
	// L3 must always be nil for 3rd-party results.
	// (VantageResult has no L3 field — this is a compile-time guarantee.)

	if createCalls.Load() != 3 {
		t.Errorf("expected 3 CreateSiteMonitor calls (one per ISP), got %d", createCalls.Load())
	}
}

// TestProbeRateLimit_DegradesToNoData verifies that a 429 response causes
// the agent to return an empty slice (no data) and NOT write any Redis keys.
func TestProbeRateLimit_DegradesToNoData(t *testing.T) {
	mr := miniredis.RunT(t)
	srv, _, _ := newMockAliyunServer(t, behaviorRateLimit)
	agent := newTestAgent(t, srv, mr)

	target := probe.ProbeTarget{NodeID: "node-sg-01", Host: "sg.example.com", Port: 443}
	results, err := agent.Probe(context.Background(), target)
	if err != nil {
		t.Fatalf("Probe should not return error on rate-limit (degraded); got: %v", err)
	}
	if len(results) != 0 {
		t.Errorf("expected 0 results on rate-limit, got %d", len(results))
	}

	// No Redis keys should have been written.
	keys := mr.Keys()
	for _, k := range keys {
		if strings.HasPrefix(k, "probe:node-sg-01") {
			t.Errorf("unexpected Redis key written during degradation: %q", k)
		}
	}
}

// TestProbeTimeout_DegradesToNoData verifies that a hung vendor server (simulating
// timeout) causes degradation: empty results, no Redis writes.
func TestProbeTimeout_DegradesToNoData(t *testing.T) {
	mr := miniredis.RunT(t)
	srv, _, _ := newMockAliyunServer(t, behaviorTimeout)

	rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
	st := probe.NewStore(rdb)
	cfg := probe.AliyunSyntheticAgentConfig{
		AccessKeyID:     "test-key-id",
		AccessKeySecret: "test-key-secret",
		PollInterval:    10 * time.Millisecond,
		PollTimeout:     100 * time.Millisecond, // very short for test
	}
	// Use a short HTTP client timeout too.
	agent := probe.NewAliyunSyntheticAgentWithOptions(cfg, st, slog.Default(), srv.URL+"/")
	_ = probe.SetHTTPClientTimeout(agent, 50*time.Millisecond)

	target := probe.ProbeTarget{NodeID: "node-jp-01", Host: "jp.example.com", Port: 443}

	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
	defer cancel()

	results, err := agent.Probe(ctx, target)
	if err != nil {
		t.Fatalf("Probe should not propagate timeout error; got: %v", err)
	}
	if len(results) != 0 {
		t.Errorf("expected 0 results on timeout, got %d", len(results))
	}

	// No Redis keys for this node.
	for _, k := range mr.Keys() {
		if strings.HasPrefix(k, "probe:node-jp-01") {
			t.Errorf("unexpected Redis key on timeout: %q", k)
		}
	}
}

// TestRunOnce_WritesRedisKeys verifies that RunOnce writes probe:{node}:3rd-* keys
// to Redis for each ISP after a successful probe cycle.
func TestRunOnce_WritesRedisKeys(t *testing.T) {
	mr := miniredis.RunT(t)
	srv, _, _ := newMockAliyunServer(t, behaviorSuccess)
	agent := newTestAgent(t, srv, mr)

	targets := []probe.ProbeTarget{
		{NodeID: "node-hk-02", Host: "hk2.example.com", Port: 443},
	}
	if err := agent.RunOnce(context.Background(), targets); err != nil {
		t.Fatalf("RunOnce: %v", err)
	}

	// Expect 3 keys: one per ISP, each with "3rd-" in the vantage segment.
	var thirdPartyKeys []string
	for _, k := range mr.Keys() {
		if strings.HasPrefix(k, "probe:node-hk-02:") && strings.Contains(k, "3rd-") {
			thirdPartyKeys = append(thirdPartyKeys, k)
		}
	}
	if len(thirdPartyKeys) != 3 {
		t.Errorf("expected 3 probe:node-hk-02:...:3rd-* keys, got %d: %v", len(thirdPartyKeys), thirdPartyKeys)
	}

	// Verify key format: probe:{node}:{country}:{region}:{isp}
	// e.g. probe:node-hk-02:CN:北京:3rd-ChinaTelecom  (spaces become underscores)
	for _, k := range thirdPartyKeys {
		parts := strings.Split(k, ":")
		// parts[0]=probe, parts[1]=node-hk-02, parts[2]=CN, parts[3]=region, parts[4]=3rd-...
		if len(parts) < 5 {
			t.Errorf("unexpected key format: %q", k)
			continue
		}
		if !strings.HasPrefix(parts[4], "3rd-") {
			t.Errorf("ISP segment %q missing 3rd- prefix in key %q", parts[4], k)
		}
	}
}

// TestProbeFailedTarget_L1NotOK verifies that when the vendor reports the target
// is unreachable (Availability=0), L1.OK is false and no L2 is populated.
func TestProbeFailedTarget_L1NotOK(t *testing.T) {
	mr := miniredis.RunT(t)
	srv, _, _ := newMockAliyunServer(t, behaviorFailedProbe)
	agent := newTestAgent(t, srv, mr)

	target := probe.ProbeTarget{NodeID: "node-fail-01", Host: "fail.example.com", Port: 443}
	results, err := agent.Probe(context.Background(), target)
	if err != nil {
		t.Fatalf("Probe: %v", err)
	}
	if len(results) != 3 {
		t.Fatalf("expected 3 results even on target failure, got %d", len(results))
	}
	for _, vr := range results {
		if vr.L1.OK {
			t.Errorf("ISP %s: L1.OK should be false for unreachable target", vr.Vantage.ISP)
		}
		if vr.L2 != nil {
			t.Errorf("ISP %s: L2 should be nil when L1 failed", vr.Vantage.ISP)
		}
	}
}

// TestNoPlaintextCredentials is a compile-time reminder that credentials
// must not be hardcoded.  The test itself only validates that the agent
// accepts empty credentials (real validation is at runtime via env vars).
func TestNoPlaintextCredentials(t *testing.T) {
	cfg := probe.AliyunSyntheticAgentConfig{
		AccessKeyID:     "", // must come from env, not code
		AccessKeySecret: "", // must come from env, not code
	}
	// A zero-credential config is valid to construct; errors surface on actual calls.
	if cfg.AccessKeyID != "" || cfg.AccessKeySecret != "" {
		t.Error("credentials must not be hardcoded in config defaults")
	}
}