387 lines
12 KiB
Go
387 lines
12 KiB
Go
package discovery
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
|
|
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
|
|
"k8s.io/apimachinery/pkg/runtime"
|
|
"k8s.io/apimachinery/pkg/types"
|
|
"sigs.k8s.io/controller-runtime/pkg/client/fake"
|
|
|
|
crawlv1alpha1 "gitea.home.hrajfrisbee.cz/kacerr/egress-proxies-operator/api/v1alpha1"
|
|
"gitea.home.hrajfrisbee.cz/kacerr/egress-proxies-operator/internal/lease"
|
|
)
|
|
|
|
func testProxy(name string, attrs map[string]string, healthy bool, mut ...func(*crawlv1alpha1.Proxy)) *crawlv1alpha1.Proxy {
|
|
p := &crawlv1alpha1.Proxy{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: name, Namespace: "default", UID: types.UID("uid-" + name),
|
|
},
|
|
Spec: crawlv1alpha1.ProxySpec{
|
|
Mode: crawlv1alpha1.ModeExternal,
|
|
Endpoint: &crawlv1alpha1.EndpointSpec{Host: "10.0.0.1", Port: 3128},
|
|
Attributes: attrs,
|
|
},
|
|
}
|
|
p.Status.IP = "10.0.0.1"
|
|
p.Status.Phase = crawlv1alpha1.PhaseReady
|
|
status := metav1.ConditionFalse
|
|
if healthy {
|
|
status = metav1.ConditionTrue
|
|
}
|
|
p.Status.Conditions = []metav1.Condition{{
|
|
Type: crawlv1alpha1.ConditionHealthy, Status: status,
|
|
Reason: "Probing", LastTransitionTime: metav1.Now(),
|
|
}}
|
|
for _, m := range mut {
|
|
m(p)
|
|
}
|
|
return p
|
|
}
|
|
|
|
func withMaxLeases(n int32) func(*crawlv1alpha1.Proxy) {
|
|
return func(p *crawlv1alpha1.Proxy) { p.Spec.MaxLeases = &n }
|
|
}
|
|
|
|
func withLatency(ms int64) func(*crawlv1alpha1.Proxy) {
|
|
return func(p *crawlv1alpha1.Proxy) { p.Status.LatencyMillis = ms }
|
|
}
|
|
|
|
// newTestServer wires the handler chain to a fake cache reader and a real
|
|
// lease store, served over httptest.
|
|
func newTestServer(t *testing.T, token string, proxies ...*crawlv1alpha1.Proxy) (*httptest.Server, *Server) {
|
|
t.Helper()
|
|
s := runtime.NewScheme()
|
|
if err := crawlv1alpha1.AddToScheme(s); err != nil {
|
|
t.Fatalf("scheme: %v", err)
|
|
}
|
|
builder := fake.NewClientBuilder().WithScheme(s)
|
|
for _, p := range proxies {
|
|
builder = builder.WithObjects(p)
|
|
}
|
|
srv := &Server{
|
|
Reader: builder.Build(),
|
|
Store: lease.NewStore(15 * time.Minute),
|
|
Token: token,
|
|
MaxLeaseTTL: time.Hour,
|
|
}
|
|
ts := httptest.NewServer(srv.handler())
|
|
t.Cleanup(ts.Close)
|
|
return ts, srv
|
|
}
|
|
|
|
type response struct {
|
|
status int
|
|
body map[string]any
|
|
}
|
|
|
|
func do(t *testing.T, ts *httptest.Server, method, path, token string, body any) response {
|
|
t.Helper()
|
|
var reader io.Reader
|
|
if body != nil {
|
|
if s, ok := body.(string); ok {
|
|
reader = bytes.NewBufferString(s)
|
|
} else {
|
|
b, err := json.Marshal(body)
|
|
if err != nil {
|
|
t.Fatalf("marshaling request body: %v", err)
|
|
}
|
|
reader = bytes.NewBuffer(b)
|
|
}
|
|
}
|
|
req, err := http.NewRequestWithContext(context.Background(), method, ts.URL+path, reader)
|
|
if err != nil {
|
|
t.Fatalf("building request: %v", err)
|
|
}
|
|
if token != "" {
|
|
req.Header.Set("Authorization", "Bearer "+token)
|
|
}
|
|
resp, err := ts.Client().Do(req)
|
|
if err != nil {
|
|
t.Fatalf("%s %s: %v", method, path, err)
|
|
}
|
|
defer func() { _ = resp.Body.Close() }()
|
|
out := response{status: resp.StatusCode}
|
|
raw, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
t.Fatalf("reading response: %v", err)
|
|
}
|
|
if len(raw) > 0 && resp.Header.Get("Content-Type") == "application/json" {
|
|
if err := json.Unmarshal(raw, &out.body); err != nil {
|
|
t.Fatalf("decoding response %q: %v", raw, err)
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
func TestAuth(t *testing.T) {
|
|
t.Parallel()
|
|
ts, _ := newTestServer(t, "sekrit", testProxy("p1", nil, true))
|
|
|
|
if got := do(t, ts, http.MethodGet, "/v1/proxies", "", nil); got.status != http.StatusUnauthorized {
|
|
t.Errorf("no token: status %d, want 401", got.status)
|
|
}
|
|
if got := do(t, ts, http.MethodGet, "/v1/proxies", "wrong", nil); got.status != http.StatusUnauthorized {
|
|
t.Errorf("wrong token: status %d, want 401", got.status)
|
|
}
|
|
if got := do(t, ts, http.MethodGet, "/v1/proxies", "sekrit", nil); got.status != http.StatusOK {
|
|
t.Errorf("correct token: status %d, want 200", got.status)
|
|
}
|
|
if got := do(t, ts, http.MethodGet, "/healthz", "", nil); got.status != http.StatusOK {
|
|
t.Errorf("healthz without token: status %d, want 200 (always unauthenticated)", got.status)
|
|
}
|
|
}
|
|
|
|
func TestAuth_disabledWithEmptyToken(t *testing.T) {
|
|
t.Parallel()
|
|
ts, _ := newTestServer(t, "", testProxy("p1", nil, true))
|
|
if got := do(t, ts, http.MethodGet, "/v1/proxies", "", nil); got.status != http.StatusOK {
|
|
t.Errorf("status %d, want 200 with auth disabled", got.status)
|
|
}
|
|
}
|
|
|
|
func TestListProxies(t *testing.T) {
|
|
t.Parallel()
|
|
ts, _ := newTestServer(t, "",
|
|
testProxy("eu-healthy", map[string]string{"geo": "eu", "purpose": "crawl"}, true),
|
|
testProxy("eu-sick", map[string]string{"geo": "eu"}, false),
|
|
testProxy("us-healthy", map[string]string{"geo": "us"}, true),
|
|
)
|
|
|
|
tests := []struct {
|
|
name string
|
|
query string
|
|
wantCount int
|
|
wantFirst string
|
|
}{
|
|
{name: "no filter returns everything", query: "", wantCount: 3, wantFirst: "default/eu-healthy"},
|
|
{name: "healthy filter", query: "?healthy=true", wantCount: 2},
|
|
{name: "unhealthy filter", query: "?healthy=false", wantCount: 1, wantFirst: "default/eu-sick"},
|
|
{name: "attribute filter", query: "?attr.geo=eu", wantCount: 2},
|
|
{name: "attribute and health combined", query: "?attr.geo=eu&healthy=true", wantCount: 1, wantFirst: "default/eu-healthy"},
|
|
{name: "two attributes must both match", query: "?attr.geo=eu&attr.purpose=crawl", wantCount: 1, wantFirst: "default/eu-healthy"},
|
|
{name: "no matches is 200 with count 0", query: "?attr.geo=mars", wantCount: 0},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
got := do(t, ts, http.MethodGet, "/v1/proxies"+tc.query, "", nil)
|
|
if got.status != http.StatusOK {
|
|
t.Fatalf("status %d, want 200", got.status)
|
|
}
|
|
count := int(got.body["count"].(float64))
|
|
proxies := got.body["proxies"].([]any)
|
|
if count != tc.wantCount || len(proxies) != tc.wantCount {
|
|
t.Fatalf("count = %d (len %d), want %d", count, len(proxies), tc.wantCount)
|
|
}
|
|
if tc.wantFirst != "" {
|
|
first := proxies[0].(map[string]any)
|
|
if first["id"] != tc.wantFirst {
|
|
t.Errorf("first id = %v, want %s", first["id"], tc.wantFirst)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
|
|
t.Run("invalid healthy value is 400", func(t *testing.T) {
|
|
t.Parallel()
|
|
if got := do(t, ts, http.MethodGet, "/v1/proxies?healthy=maybe", "", nil); got.status != http.StatusBadRequest {
|
|
t.Errorf("status %d, want 400", got.status)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestAcquireLease_grantShape(t *testing.T) {
|
|
t.Parallel()
|
|
ts, _ := newTestServer(t, "",
|
|
testProxy("eu1", map[string]string{"geo": "eu"}, true, withLatency(30)),
|
|
testProxy("eu2", map[string]string{"geo": "eu"}, true, withLatency(10)),
|
|
)
|
|
|
|
got := do(t, ts, http.MethodPost, "/v1/leases", "", map[string]any{
|
|
"selector": map[string]string{"geo": "eu"},
|
|
})
|
|
if got.status != http.StatusCreated {
|
|
t.Fatalf("status %d (%v), want 201", got.status, got.body)
|
|
}
|
|
if got.body["leaseID"] == "" || got.body["leaseID"] == nil {
|
|
t.Error("empty leaseID")
|
|
}
|
|
if got.body["ttlSeconds"].(float64) != 300 {
|
|
t.Errorf("ttlSeconds = %v, want the 300 default", got.body["ttlSeconds"])
|
|
}
|
|
proxy := got.body["proxy"].(map[string]any)
|
|
if proxy["id"] != "default/eu2" {
|
|
t.Errorf("granted %v, want default/eu2 (lower latency at equal load)", proxy["id"])
|
|
}
|
|
if proxy["activeLeases"].(float64) != 1 {
|
|
t.Errorf("activeLeases = %v, want 1 (this grant included)", proxy["activeLeases"])
|
|
}
|
|
if _, err := time.Parse(time.RFC3339, got.body["expiresAt"].(string)); err != nil {
|
|
t.Errorf("expiresAt %v is not RFC3339: %v", got.body["expiresAt"], err)
|
|
}
|
|
}
|
|
|
|
func TestAcquireLease_noMatchBody(t *testing.T) {
|
|
t.Parallel()
|
|
ts, _ := newTestServer(t, "",
|
|
testProxy("eu-tiny", map[string]string{"geo": "eu"}, true, withMaxLeases(1)),
|
|
testProxy("eu-sick", map[string]string{"geo": "eu"}, false),
|
|
)
|
|
|
|
body := map[string]any{"selector": map[string]string{"geo": "eu"}}
|
|
if got := do(t, ts, http.MethodPost, "/v1/leases", "", body); got.status != http.StatusCreated {
|
|
t.Fatalf("first acquire: status %d, want 201", got.status)
|
|
}
|
|
|
|
got := do(t, ts, http.MethodPost, "/v1/leases", "", body)
|
|
if got.status != http.StatusConflict {
|
|
t.Fatalf("second acquire: status %d, want 409", got.status)
|
|
}
|
|
want := map[string]float64{"considered": 2, "atCapacity": 1, "inCooldown": 0, "unhealthy": 1}
|
|
for k, v := range want {
|
|
if got.body[k].(float64) != v {
|
|
t.Errorf("%s = %v, want %v (body %v)", k, got.body[k], v, got.body)
|
|
}
|
|
}
|
|
if got.body["error"] != "no_match" {
|
|
t.Errorf("error = %v, want no_match", got.body["error"])
|
|
}
|
|
}
|
|
|
|
func TestAcquireLease_badRequests(t *testing.T) {
|
|
t.Parallel()
|
|
ts, _ := newTestServer(t, "", testProxy("p1", nil, true))
|
|
|
|
tests := []struct {
|
|
name string
|
|
body any
|
|
wantCode string
|
|
}{
|
|
{name: "ttl above the cap", body: map[string]any{"ttlSeconds": 999999}, wantCode: "invalid_ttl"},
|
|
{name: "negative ttl", body: map[string]any{"ttlSeconds": -5}, wantCode: "invalid_ttl"},
|
|
{name: "malformed json", body: "{not json", wantCode: "invalid_body"},
|
|
}
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
got := do(t, ts, http.MethodPost, "/v1/leases", "", tc.body)
|
|
if got.status != http.StatusBadRequest || got.body["error"] != tc.wantCode {
|
|
t.Errorf("= %d/%v, want 400/%s", got.status, got.body["error"], tc.wantCode)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestReleaseLease_alwaysNoContent(t *testing.T) {
|
|
t.Parallel()
|
|
ts, _ := newTestServer(t, "", testProxy("p1", nil, true))
|
|
|
|
got := do(t, ts, http.MethodPost, "/v1/leases", "", map[string]any{})
|
|
if got.status != http.StatusCreated {
|
|
t.Fatalf("acquire: status %d, want 201", got.status)
|
|
}
|
|
id := got.body["leaseID"].(string)
|
|
|
|
for _, path := range []string{"/v1/leases/" + id, "/v1/leases/" + id, "/v1/leases/never-existed"} {
|
|
if got := do(t, ts, http.MethodDelete, path, "", nil); got.status != http.StatusNoContent {
|
|
t.Errorf("DELETE %s: status %d, want 204", path, got.status)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestReportLease(t *testing.T) {
|
|
t.Parallel()
|
|
ts, _ := newTestServer(t, "", testProxy("p1", map[string]string{"geo": "eu"}, true))
|
|
|
|
got := do(t, ts, http.MethodPost, "/v1/leases", "", map[string]any{
|
|
"selector": map[string]string{"geo": "eu"}, "target": "example.com",
|
|
})
|
|
if got.status != http.StatusCreated {
|
|
t.Fatalf("acquire: status %d, want 201", got.status)
|
|
}
|
|
id := got.body["leaseID"].(string)
|
|
reportPath := fmt.Sprintf("/v1/leases/%s/report", id)
|
|
|
|
if got := do(t, ts, http.MethodPost, reportPath, "", map[string]any{"result": "rate_limited", "target": "example.com"}); got.status != http.StatusNoContent {
|
|
t.Fatalf("report: status %d, want 204", got.status)
|
|
}
|
|
|
|
// The cooldown from the report now blocks same-target acquisition.
|
|
got = do(t, ts, http.MethodPost, "/v1/leases", "", map[string]any{
|
|
"selector": map[string]string{"geo": "eu"}, "target": "example.com",
|
|
})
|
|
if got.status != http.StatusConflict || got.body["inCooldown"].(float64) != 1 {
|
|
t.Errorf("post-report acquire = %d/%v, want 409 with inCooldown 1", got.status, got.body)
|
|
}
|
|
|
|
t.Run("invalid result value", func(t *testing.T) {
|
|
got := do(t, ts, http.MethodPost, reportPath, "", map[string]any{"result": "throttled"})
|
|
if got.status != http.StatusBadRequest || got.body["error"] != "invalid_result" {
|
|
t.Errorf("= %d/%v, want 400/invalid_result", got.status, got.body["error"])
|
|
}
|
|
})
|
|
t.Run("unknown lease", func(t *testing.T) {
|
|
got := do(t, ts, http.MethodPost, "/v1/leases/never-existed/report", "", map[string]any{"result": "ok"})
|
|
if got.status != http.StatusNotFound || got.body["error"] != "unknown_lease" {
|
|
t.Errorf("= %d/%v, want 404/unknown_lease", got.status, got.body["error"])
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestStart_servesAndShutsDown(t *testing.T) {
|
|
t.Parallel()
|
|
s := runtime.NewScheme()
|
|
if err := crawlv1alpha1.AddToScheme(s); err != nil {
|
|
t.Fatalf("scheme: %v", err)
|
|
}
|
|
srv := &Server{
|
|
Reader: fake.NewClientBuilder().WithScheme(s).Build(),
|
|
Store: lease.NewStore(time.Minute),
|
|
Addr: "127.0.0.1:0",
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
done := make(chan error, 1)
|
|
go func() { done <- srv.Start(ctx) }()
|
|
|
|
var addr string
|
|
deadline := time.After(5 * time.Second)
|
|
for addr == "" {
|
|
select {
|
|
case <-deadline:
|
|
t.Fatal("server never bound")
|
|
case <-time.After(5 * time.Millisecond):
|
|
addr = srv.BoundAddr()
|
|
}
|
|
}
|
|
|
|
resp, err := http.Get("http://" + addr + "/healthz")
|
|
if err != nil {
|
|
t.Fatalf("healthz: %v", err)
|
|
}
|
|
_ = resp.Body.Close()
|
|
if resp.StatusCode != http.StatusOK {
|
|
t.Errorf("healthz status %d, want 200", resp.StatusCode)
|
|
}
|
|
|
|
cancel()
|
|
select {
|
|
case err := <-done:
|
|
if err != nil {
|
|
t.Errorf("Start returned %v, want nil after graceful shutdown", err)
|
|
}
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("Start did not stop on cancel")
|
|
}
|
|
}
|