whoissdk/rdap_client_retry_test.go
2026-03-19 11:53:07 +08:00

95 lines
2.5 KiB
Go

package whois
import (
"context"
"net/http"
"net/http/httptest"
"sync/atomic"
"testing"
"time"
)
func TestRDAPClientRetryOn5xx(t *testing.T) {
var hit int32
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
n := atomic.AddInt32(&hit, 1)
if n == 1 {
http.Error(w, "temporary", http.StatusInternalServerError)
return
}
_, _ = w.Write([]byte(`{"objectClassName":"domain","ldhName":"example.com"}`))
}))
defer srv.Close()
c, err := NewRDAPClientWithBootstrap(&RDAPBootstrap{
Version: "1.0",
Services: []RDAPService{
{TLDs: []string{"com"}, URLs: []string{srv.URL}},
},
})
if err != nil {
t.Fatalf("NewRDAPClientWithBootstrap() error: %v", err)
}
c.SetRetryPolicy(RDAPRetryPolicy{
MaxAttempts: 2,
BaseDelay: time.Millisecond,
MaxDelay: 5 * time.Millisecond,
RetryOn429: true,
RetryOn5xx: true,
RetryOnNetwork: true,
})
resp, err := c.QueryDomain(context.Background(), "example.com")
if err != nil {
t.Fatalf("QueryDomain() error: %v", err)
}
if resp.StatusCode != http.StatusOK {
t.Fatalf("unexpected status: %d", resp.StatusCode)
}
if got := atomic.LoadInt32(&hit); got != 2 {
t.Fatalf("expected retry hit=2, got=%d", got)
}
}
func TestRDAPClientQueryIPAndASN(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/ip/1.1.1.1":
_, _ = w.Write([]byte(`{"objectClassName":"ip network","handle":"NET-1-1-1-0-1"}`))
case "/autnum/13335":
_, _ = w.Write([]byte(`{"objectClassName":"autnum","handle":"AS13335"}`))
default:
http.NotFound(w, r)
}
}))
defer srv.Close()
c, err := NewRDAPClientWithBootstrap(&RDAPBootstrap{
Version: "1.0",
Services: []RDAPService{
{TLDs: []string{"com"}, URLs: []string{"https://rdap.example.com/"}},
},
})
if err != nil {
t.Fatalf("NewRDAPClientWithBootstrap() error: %v", err)
}
c.SetIPServers(srv.URL)
c.SetASNServers(srv.URL)
ipResp, err := c.QueryIP(context.Background(), "1.1.1.1")
if err != nil {
t.Fatalf("QueryIP() error: %v", err)
}
if ipResp.StatusCode != http.StatusOK || ipResp.Domain != "1.1.1.1" {
t.Fatalf("unexpected ip response: status=%d domain=%q", ipResp.StatusCode, ipResp.Domain)
}
asnResp, err := c.QueryASN(context.Background(), "AS13335")
if err != nil {
t.Fatalf("QueryASN() error: %v", err)
}
if asnResp.StatusCode != http.StatusOK || asnResp.Domain != "AS13335" {
t.Fatalf("unexpected asn response: status=%d domain=%q", asnResp.StatusCode, asnResp.Domain)
}
}