EvoBGP/internal/pipeline/doh_resolve_test.go

package pipeline

import (
	"context"
	"net/http"
	"net/http/httptest"
	"testing"

	"evobgp/internal/store"
)

func TestResolveDomainIPsWithPolicy_Union(t *testing.T) {
	srvRU := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte(`{"Answer":[{"type":1,"data":"198.51.100.1"}]}`))
	}))
	defer srvRU.Close()
	srvEU := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte(`{"Answer":[{"type":1,"data":"203.0.113.1"}]}`))
	}))
	defer srvEU.Close()

	profiles := []*store.DohProfile{
		{URL: srvRU.URL},
		{URL: srvEU.URL},
	}
	ips, err := resolveDomainIPsWithPolicy(context.Background(), srvRU.Client(), profiles, store.DohPolicyUnion, "example.com")
	if err != nil {
		t.Fatal(err)
	}
	if len(ips) != 2 {
		t.Fatalf("want 2 ips, got %v", ips)
	}
	seen := map[string]bool{ips[0].String(): true, ips[1].String(): true}
	if !seen["198.51.100.1"] || !seen["203.0.113.1"] {
		t.Fatalf("unexpected ips: %v", ips)
	}
}

func TestResolveDomainIPsWithPolicy_Failover(t *testing.T) {
	var calls int
	srvBad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls++
		http.Error(w, "fail", http.StatusBadGateway)
	}))
	defer srvBad.Close()
	srvOK := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls++
		_, _ = w.Write([]byte(`{"Answer":[{"type":1,"data":"198.51.100.5"}]}`))
	}))
	defer srvOK.Close()

	profiles := []*store.DohProfile{
		{URL: srvBad.URL},
		{URL: srvOK.URL},
	}
	ips, err := resolveDomainIPsWithPolicy(context.Background(), srvBad.Client(), profiles, store.DohPolicyFailover, "example.com")
	if err != nil {
		t.Fatal(err)
	}
	if len(ips) != 1 || ips[0].String() != "198.51.100.5" {
		t.Fatalf("unexpected ips: %v", ips)
	}
	if calls < 2 {
		t.Fatalf("want at least 2 resolver calls, got %d", calls)
	}
}

func TestResolveDomainIPsWithPolicy_PrimaryOnly(t *testing.T) {
	var secondCalled bool
	srv1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte(`{"Answer":[{"type":1,"data":"198.51.100.9"}]}`))
	}))
	defer srv1.Close()
	srv2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		secondCalled = true
		_, _ = w.Write([]byte(`{"Answer":[{"type":1,"data":"203.0.113.9"}]}`))
	}))
	defer srv2.Close()

	profiles := []*store.DohProfile{
		{URL: srv1.URL},
		{URL: srv2.URL},
	}
	ips, err := resolveDomainIPsWithPolicy(context.Background(), srv1.Client(), profiles, store.DohPolicyPrimaryOnly, "example.com")
	if err != nil {
		t.Fatal(err)
	}
	if len(ips) != 1 || ips[0].String() != "198.51.100.9" {
		t.Fatalf("unexpected ips: %v", ips)
	}
	if secondCalled {
		t.Fatal("secondary resolver must not be queried in primary_only mode")
	}
}

func TestCollectModulePrefixRows_DohUnion(t *testing.T) {
	m := store.NewMemory()
	m.SeedDemo()
	tenant, _, _, _, _ := m.DemoIDs()

	mod, err := m.CreateModule(tenant, &store.Module{
		Type:              "DOMAINS",
		Name:              "domains-union",
		Enabled:           true,
		DohResolverPolicy: store.DohPolicyUnion,
	})
	if err != nil {
		t.Fatal(err)
	}

	srvRU := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte(`{"Answer":[{"type":1,"data":"198.51.100.2"}]}`))
	}))
	defer srvRU.Close()
	srvEU := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte(`{"Answer":[{"type":1,"data":"203.0.113.2"}]}`))
	}))
	defer srvEU.Close()

	ru, err := m.CreateDohProfile(tenant, &store.DohProfile{Name: "ru", URL: srvRU.URL})
	if err != nil {
		t.Fatal(err)
	}
	eu, err := m.CreateDohProfile(tenant, &store.DohProfile{Name: "eu", URL: srvEU.URL})
	if err != nil {
		t.Fatal(err)
	}
	if _, err := m.UpdateModule(tenant, mod.ID, &store.ModulePatch{
		DohProfileIDs: &[]string{ru.ID, eu.ID},
	}); err != nil {
		t.Fatal(err)
	}
	mod, err = m.GetModule(tenant, mod.ID)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := m.CreateDomainEntry(tenant, mod.ID, &store.DomainEntry{FQDN: "svc.example.com"}); err != nil {
		t.Fatal(err)
	}

	rows, err := collectModulePrefixRows(context.Background(), m, srvRU.Client(), tenant, mod, nil)
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) != 2 {
		t.Fatalf("want 2 prefix rows, got %+v", rows)
	}
}