Privatae
106 lines
3.5 KiB
Go
106 lines
3.5 KiB
Go
package main
|
|
|
|
import (
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
// piiProbe is opt-in: no PRIVATAE_PII_ADDR means no probe, no capability, no
|
|
// heartbeat fields — a box without the service is unaffected by this lane.
|
|
func TestPiiProbeRequiresEnv(t *testing.T) {
|
|
t.Setenv("PRIVATAE_PII_ADDR", "")
|
|
if adv, ok := piiProbe(); adv || ok {
|
|
t.Errorf("piiProbe(no env) = (%v, %v), want (false, false)", adv, ok)
|
|
}
|
|
}
|
|
|
|
func TestPiiProbeHealthy(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/health" {
|
|
t.Errorf("probe hit %s, want /health", r.URL.Path)
|
|
}
|
|
w.Write([]byte(`{"ok":true}`))
|
|
}))
|
|
defer srv.Close()
|
|
t.Setenv("PRIVATAE_PII_ADDR", strings.TrimPrefix(srv.URL, "http://"))
|
|
if adv, ok := piiProbe(); !adv || !ok {
|
|
t.Errorf("piiProbe(healthy) = (%v, %v), want (true, true)", adv, ok)
|
|
}
|
|
}
|
|
|
|
// Env set but upstream dead: still advertised (the forwarder runs), with
|
|
// pii_ok=false as the distinct dead-upstream signal for the host.
|
|
func TestPiiProbeDeadUpstreamStillAdvertised(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusBadGateway)
|
|
}))
|
|
addr := strings.TrimPrefix(srv.URL, "http://")
|
|
srv.Close() // connection refused from here on
|
|
t.Setenv("PRIVATAE_PII_ADDR", addr)
|
|
if adv, ok := piiProbe(); !adv || ok {
|
|
t.Errorf("piiProbe(dead upstream) = (%v, %v), want (true, false)", adv, ok)
|
|
}
|
|
}
|
|
|
|
func TestHeartbeatAdvertisesPii(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Write([]byte(`{"ok":true}`))
|
|
}))
|
|
defer srv.Close()
|
|
t.Setenv("PRIVATAE_PII_ADDR", strings.TrimPrefix(srv.URL, "http://"))
|
|
t.Setenv("PRIVATAE_LLM_ADDR", "")
|
|
|
|
raw := buildHeartbeatBody("", "h", "linux", "test", "pk", "fp", "100.64.0.9", "ep", false, nil, nil, nil)
|
|
var body struct {
|
|
Capabilities []string `json:"capabilities"`
|
|
PiiPort *int `json:"pii_port"`
|
|
PiiOK *bool `json:"pii_ok"`
|
|
}
|
|
if err := json.Unmarshal([]byte(raw), &body); err != nil {
|
|
t.Fatalf("heartbeat body: %v", err)
|
|
}
|
|
found := false
|
|
for _, c := range body.Capabilities {
|
|
if c == "pii" {
|
|
found = true
|
|
}
|
|
}
|
|
if !found {
|
|
t.Errorf("capabilities = %v, want to include \"pii\"", body.Capabilities)
|
|
}
|
|
if body.PiiPort == nil || *body.PiiPort != piiForwardPort {
|
|
t.Errorf("pii_port = %v, want %d (the port servePIIForward listens on)", body.PiiPort, piiForwardPort)
|
|
}
|
|
if body.PiiOK == nil || !*body.PiiOK {
|
|
t.Errorf("pii_ok = %v, want true", body.PiiOK)
|
|
}
|
|
}
|
|
|
|
// Absence semantics: unset env produces NO pii fields at all — the host
|
|
// reads a missing pii_port as "service not present", so an always-present
|
|
// field (like llm_port's style) would break the honest disabled state.
|
|
func TestHeartbeatOmitsPiiWhenUnset(t *testing.T) {
|
|
t.Setenv("PRIVATAE_PII_ADDR", "")
|
|
t.Setenv("PRIVATAE_LLM_ADDR", "")
|
|
raw := buildHeartbeatBody("", "h", "linux", "test", "pk", "fp", "100.64.0.9", "ep", false, nil, nil, nil)
|
|
var body map[string]interface{}
|
|
if err := json.Unmarshal([]byte(raw), &body); err != nil {
|
|
t.Fatalf("heartbeat body: %v", err)
|
|
}
|
|
if _, present := body["pii_port"]; present {
|
|
t.Error("pii_port present with PRIVATAE_PII_ADDR unset; want omitted")
|
|
}
|
|
if _, present := body["pii_ok"]; present {
|
|
t.Error("pii_ok present with PRIVATAE_PII_ADDR unset; want omitted")
|
|
}
|
|
caps, _ := body["capabilities"].([]interface{})
|
|
for _, c := range caps {
|
|
if c == "pii" {
|
|
t.Error("capability \"pii\" present with PRIVATAE_PII_ADDR unset")
|
|
}
|
|
}
|
|
}
|