selfheal-lab-apps/ai/internal/agent/agent_test.go

package agent

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/diag"
	"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/harness"
	"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/openai"
	"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/tools"
)

// A mock Responses API: turn 1 asks for get_probe_status, turn 2 submits the diagnosis. Verifies the loop wiring
// (function_call_output with the call_id, previous_response_id chaining, strict submit parsing).
func TestDiagnoseLoop(t *testing.T) {
	var turns int
	var sawOutput, sawPrev bool
	oaiSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/embeddings" {
			w.Write([]byte(`{"data":[{"index":0,"embedding":[0.1,0.2]}]}`))
			return
		}
		var req openai.Request
		_ = json.NewDecoder(r.Body).Decode(&req)
		turns++
		for _, it := range req.Input {
			if it.Type == "function_call_output" && it.CallID == "call_1" && strings.Contains(it.Output, "probe_success") {
				sawOutput = true
			}
		}
		if req.PreviousResponseID == "resp_1" {
			sawPrev = true
		}
		if r.Header.Get("Authorization") != "Bearer test-key" {
			t.Errorf("missing bearer")
		}
		switch turns {
		case 1:
			w.Write([]byte(`{"id":"resp_1","status":"completed","output":[{"type":"function_call","call_id":"call_1","name":"get_probe_status","arguments":"{}"}],"usage":{"input_tokens":10,"output_tokens":5}}`))
		default:
			args := `{\"root_cause_label\":\"db_primary_down\",\"confidence\":0.8,\"evidence\":[\"probe_success db-primary=0\"],\"proposed_action\":\"escalate\",\"target\":\"db-primary\",\"rationale\":\"primary unreachable\"}`
			w.Write([]byte(`{"id":"resp_2","status":"completed","output":[{"type":"function_call","call_id":"call_2","name":"submit_diagnosis","arguments":"` + args + `"}],"usage":{"input_tokens":20,"output_tokens":9}}`))
		}
	}))
	defer oaiSrv.Close()
	mimir := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasSuffix(r.URL.Path, "/api/v1/query") {
			w.Write([]byte(`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"target":"db-primary"},"value":[1,"0"]},{"metric":{"target":"alb"},"value":[1,"1"]}]}}`))
			return
		}
		w.Write([]byte(`{"status":"success","data":{"resultType":"matrix","result":[]}}`))
	}))
	defer mimir.Close()
	oai := openai.New(oaiSrv.URL, func() (string, error) { return "test-key", nil })
	tr := tools.New(tools.Config{MimirURL: mimir.URL, LokiURL: mimir.URL, EmbedModel: "e"}, nil, oai)
	enums := diag.Enums{Labels: diag.DefaultLabels, Actions: diag.DefaultActions}
	ag := New(Config{Model: "m", SystemPrompt: "sys", Topology: "topo", SLO: "slo", HitlURL: "http://127.0.0.1:1", MaxTurns: 5}, oai, tr, harness.New("http://127.0.0.1:1"), nil, enums, "pv1")
	res, err := ag.Diagnose(context.Background(), Request{Alert: map[string]any{"alertname": "DbPrimaryDown"}})
	if err != nil {
		t.Fatal(err)
	}
	if res.RootCauseLabel != "db_primary_down" || res.Turns != 2 || !sawOutput || !sawPrev {
		t.Fatalf("unexpected result: %+v sawOutput=%v sawPrev=%v", res, sawOutput, sawPrev)
	}
	ag.Dispose(context.Background(), Request{Alert: map[string]any{}}, res)
	if res.Disposition != "escalated" {
		t.Fatalf("escalate action must escalate, got %s", res.Disposition)
	}
}

func TestTargetFor(t *testing.T) {
	m := targetFor("k8s.rollout_restart", "demo/api")
	if m["namespace"] != "demo" || m["deployment"] != "api" {
		t.Fatalf("%v", m)
	}
	if targetFor("mig.resize", "")["mig"] != "selfheal-legacy-mig" {
		t.Fatal("mig")
	}
}