sub-store/internal/handler/egress_background_test.go

package handler

import (
	"context"
	"testing"
	"time"

	"github.com/peterqiu0516/sub-store/internal/model"
	"github.com/peterqiu0516/sub-store/internal/proxy"
	"github.com/peterqiu0516/sub-store/internal/service"
)

// egressKeyForNode parses a single proxy URI and returns its egress cache key
// (via the service-layer key shared with the pipeline).
func egressKeyForNode(t *testing.T, uri string) string {
	t.Helper()
	nodes := proxy.ParseProxies(uri)
	if len(nodes) != 1 {
		t.Fatalf("parse %q: got %d nodes", uri, len(nodes))
	}
	return service.EgressCacheKey(nodes[0])
}

// stubProbeNodeEgress replaces the mihomo-based probe seam for the duration
// of the test and records every probed node name.
func stubProbeNodeEgress(t *testing.T) *[]string {
	t.Helper()
	var probed []string
	old := probeNodeEgressFn
	probeNodeEgressFn = func(d *Deps, node model.ProxyNode) (map[string]any, error) {
		probed = append(probed, node["name"].(string))
		return map[string]any{"egressIp": "2.2.2.2", "country": "United States", "countryCode": "US", "flag": "🇺🇸"}, nil
	}
	t.Cleanup(func() { probeNodeEgressFn = old })
	return &probed
}

func TestRefreshAllSourceEgress_ProbesMissingSkipsCached(t *testing.T) {
	deps := newTestDeps(t)
	content := "ss://YWVzLTI1Ni1nY206cGFzcw==@1.2.3.4:80#cached-node\nss://YWVzLTI1Ni1nY206cGFzcw==@1.2.3.5:81#missing-node"
	if _, err := deps.SourceRepo.Upsert(model.SourceRecord{
		ID: "s1", Name: "s1", Type: "local", Content: content, Enabled: true,
		Filters: []model.FilterRule{}, Meta: map[string]any{},
	}); err != nil {
		t.Fatalf("upsert source: %v", err)
	}
	// Disabled sources are never refreshed.
	if _, err := deps.SourceRepo.Upsert(model.SourceRecord{
		ID: "s2", Name: "s2", Type: "local", Content: "ss://YWVzLTI1Ni1nY206cGFzcw==@1.2.3.6:82#disabled-node", Enabled: false,
		Filters: []model.FilterRule{}, Meta: map[string]any{},
	}); err != nil {
		t.Fatalf("upsert disabled source: %v", err)
	}

	// Fresh cache entry (geo result) for the first node only.
	deps.CacheRepo.SafePut(egressKeyForNode(t, "ss://YWVzLTI1Ni1nY206cGFzcw==@1.2.3.4:80#cached-node"),
		`{"egressIp":"1.1.1.1","country":"Japan","countryCode":"JP","flag":"🇯🇵"}`, nil, 300)

	probed := stubProbeNodeEgress(t)
	deps.refreshAllSourceEgress(context.Background())

	if len(*probed) != 1 || (*probed)[0] != "missing-node" {
		t.Fatalf("probed = %v, want exactly [missing-node] (cached and disabled skipped)", *probed)
	}
	// The fresh probe result must now be cached.
	if _, ok := deps.CacheRepo.SafeGet(egressKeyForNode(t, "ss://YWVzLTI1Ni1nY206cGFzcw==@1.2.3.5:81#missing-node")); !ok {
		t.Fatal("probe result should be cached after refresh")
	}
}

func TestRefreshAllSourceEgress_SkipsErrorResultsWithinTTL(t *testing.T) {
	deps := newTestDeps(t)
	if _, err := deps.SourceRepo.Upsert(model.SourceRecord{
		ID: "s1", Name: "s1", Type: "local", Content: "ss://YWVzLTI1Ni1nY206cGFzcw==@1.2.3.4:80#err-node", Enabled: true,
		Filters: []model.FilterRule{}, Meta: map[string]any{},
	}); err != nil {
		t.Fatalf("upsert source: %v", err)
	}
	deps.CacheRepo.SafePut(egressKeyForNode(t, "ss://YWVzLTI1Ni1nY206cGFzcw==@1.2.3.4:80#err-node"),
		`{"egressError":"dial tcp: timeout"}`, nil, 300)

	probed := stubProbeNodeEgress(t)
	deps.refreshAllSourceEgress(context.Background())

	if len(*probed) != 0 {
		t.Fatalf("cached error result within TTL must be skipped, probed = %v", *probed)
	}
}

func TestRefreshAllSourceEgress_ReprobesExpiredEntries(t *testing.T) {
	deps := newTestDeps(t)
	if _, err := deps.SourceRepo.Upsert(model.SourceRecord{
		ID: "s1", Name: "s1", Type: "local", Content: "ss://YWVzLTI1Ni1nY206cGFzcw==@1.2.3.4:80#stale-node", Enabled: true,
		Filters: []model.FilterRule{}, Meta: map[string]any{},
	}); err != nil {
		t.Fatalf("upsert source: %v", err)
	}
	deps.CacheRepo.SafePut(egressKeyForNode(t, "ss://YWVzLTI1Ni1nY206cGFzcw==@1.2.3.4:80#stale-node"),
		`{"egressError":"old failure"}`, nil, 300)
	// Expire every entry: SafeGet must treat them as misses.
	if _, err := deps.DB.Exec("UPDATE source_cache SET cached_at = cached_at - 10000"); err != nil {
		t.Fatalf("expire cache entries: %v", err)
	}

	probed := stubProbeNodeEgress(t)
	deps.refreshAllSourceEgress(context.Background())

	if len(*probed) != 1 {
		t.Fatalf("expired entry must be re-probed, probed = %v", *probed)
	}
	if _, ok := deps.CacheRepo.SafeGet(egressKeyForNode(t, "ss://YWVzLTI1Ni1nY206cGFzcw==@1.2.3.4:80#stale-node")); !ok {
		t.Fatal("fresh result should be re-cached after re-probe")
	}
}

// StartEgressRefresher must run the first refresh immediately and stop when
// the context is cancelled.
func TestStartEgressRefresher_ImmediateRunAndCancel(t *testing.T) {
	deps := newTestDeps(t)
	if _, err := deps.SourceRepo.Upsert(model.SourceRecord{
		ID: "s1", Name: "s1", Type: "local", Content: "ss://YWVzLTI1Ni1nY206cGFzcw==@1.2.3.4:80#n1", Enabled: true,
		Filters: []model.FilterRule{}, Meta: map[string]any{},
	}); err != nil {
		t.Fatalf("upsert source: %v", err)
	}

	done := make(chan string, 8)
	old := probeNodeEgressFn
	probeNodeEgressFn = func(d *Deps, node model.ProxyNode) (map[string]any, error) {
		done <- node["name"].(string)
		return map[string]any{"country": "United States"}, nil
	}
	t.Cleanup(func() { probeNodeEgressFn = old })

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	StartEgressRefresher(ctx, deps, time.Hour)

	select {
	case name := <-done:
		if name != "n1" {
			t.Fatalf("probed %q, want n1", name)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("refresher did not run its first pass immediately")
	}
	cancel()
}