audita/examples/library-adoption/workflow_test.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"strings"
	"sync"
	"testing"

	ck "gitea.maximumdirect.net/eric/ConfigKit"
	pk "gitea.maximumdirect.net/eric/promptkit"
)

type step struct {
	status  int
	content string
	usage   any
}

func testRuntime(t *testing.T, transport http.RoundTripper, overrides ...ck.Override) *Runtime {
	t.Helper()
	t.Setenv("AUDITA_PROTOTYPE_KEY", "synthetic-test-secret")
	loaded, err := loadConfig(ck.File("config.yaml"), overrides...)
	if err != nil {
		t.Fatal(err)
	}
	r, err := newRuntime(loaded.Config, ".", &http.Client{Transport: transport})
	if err != nil {
		t.Fatal(err)
	}
	return r
}
func intOverride(field string, value int) ck.Override {
	return ck.Override{Path: ck.Keys("llm", field), Value: value, Flag: "--" + field}
}

func TestWorkflowBoundaries(t *testing.T) {
	zero := map[string]any{"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
	for _, tc := range []struct {
		name                              string
		steps                             []step
		retries, repairs                  int
		accepted, rejected, warnings      int
		attempts, retryCount, repairCount uint64
		cause                             error
	}{
		{name: "repair both roles", steps: []step{{200, "bad", zero}, {200, validProposal, zero}, {200, "bad", nil}, {200, validAssessment, zero}}, repairs: 2, accepted: 1, attempts: 4, repairCount: 2},
		{name: "proposal repair exhausted", steps: []step{{200, "bad", nil}, {200, "bad", nil}, {200, "bad", nil}}, repairs: 2, warnings: 1, attempts: 3, repairCount: 2},
		{name: "assessment repair exhausted", steps: []step{{200, validProposal, nil}, {200, "bad", nil}, {200, "bad", nil}}, repairs: 1, rejected: 1, warnings: 1, attempts: 3, repairCount: 1},
		{name: "domain rejects missing span", steps: []step{{200, strings.Replace(validProposal, "hello world", "absent span", 1), nil}}, repairs: 2, rejected: 1, attempts: 1},
		{name: "domain rejects missing decision", steps: []step{{200, validProposal, nil}, {200, `{"validations":[]}`, nil}}, repairs: 2, rejected: 1, warnings: 1, attempts: 2},
		{name: "domain rejects duplicate decision", steps: []step{{200, validProposal, nil}, {200, `{"validations":[{"correction_index":0,"approved":true,"confidence":1,"reason":"ok"},{"correction_index":0,"approved":true,"confidence":1,"reason":"ok"}]}`, nil}}, repairs: 2, rejected: 1, warnings: 1, attempts: 2},
		{name: "assessment semantic rejection", steps: []step{{200, validProposal, nil}, {200, strings.Replace(validAssessment, `"approved":true`, `"approved":false`, 1), nil}}, repairs: 2, rejected: 1, attempts: 2},
		{name: "429 recovery", steps: []step{{429, "", nil}, {200, validProposal, zero}, {200, validAssessment, zero}}, retries: 1, accepted: 1, attempts: 3, retryCount: 1},
		{name: "explicit zero retries", steps: []step{{503, "", nil}}, repairs: 2, attempts: 1, cause: pk.ErrLLMGenerate},
		{name: "retry exhaustion after completed proposal", steps: []step{{200, validProposal, zero}, {503, "", nil}, {503, "", nil}}, retries: 1, repairs: 2, attempts: 3, retryCount: 1, cause: pk.ErrRetriesExhausted},
		{name: "retry and repair budgets", steps: []step{{503, "", nil}, {200, "bad", zero}, {503, "", nil}, {200, validProposal, zero}, {200, validAssessment, zero}}, retries: 1, repairs: 1, accepted: 1, attempts: 5, retryCount: 2, repairCount: 1},
	} {
		t.Run(tc.name, func(t *testing.T) {
			var mu sync.Mutex
			count := 0
			transport := roundTripFunc(func(req *http.Request) (*http.Response, error) {
				mu.Lock()
				defer mu.Unlock()
				if req.Header.Get("Authorization") != "Bearer synthetic-test-secret" {
					t.Error("wrong credential")
				}
				var body struct {
					Messages []struct {
						Content string `json:"content"`
					} `json:"messages"`
				}
				if err := json.NewDecoder(req.Body).Decode(&body); err != nil {
					return nil, err
				}
				if len(body.Messages) < 2 || !strings.Contains(body.Messages[1].Content, "Prompt hardening policy") {
					t.Error("real prompt hardening missing")
				}
				if count >= len(tc.steps) {
					t.Error("unexpected extra provider call")
					return providerResponse(400, "", nil), nil
				}
				s := tc.steps[count]
				count++
				return providerResponse(s.status, s.content, s.usage), nil
			})
			r := testRuntime(t, transport, intOverride("transient_retries", tc.retries), intOverride("repair_attempts", tc.repairs))
			recorder := &Recorder{}
			out, err := r.workflow(context.Background(), recorder.observer)
			if tc.cause != nil {
				if !errors.Is(err, tc.cause) {
					t.Fatalf("error %v does not match %v", err, tc.cause)
				}
			} else if err != nil {
				t.Fatal(err)
			}
			if out.Accepted != tc.accepted || out.Rejected != tc.rejected || len(out.Warnings) != tc.warnings {
				t.Fatalf("result %+v", out)
			}
			a, rt, rp := recorder.totals()
			if a != tc.attempts || rt != tc.retryCount || rp != tc.repairCount || count != len(tc.steps) {
				t.Fatalf("counts %d/%d/%d; dispatches %d", a, rt, rp, count)
			}
			records := recorder.snapshot()
			serialized, _ := json.Marshal(records)
			if strings.Contains(string(serialized), "synthetic-test-secret") || strings.Contains(string(serialized), "hello world") {
				t.Fatal("sensitive observation")
			}
			for _, o := range records {
				if o.Kind == string(pk.ExecutionAttempt) && o.StatusCode == 503 && o.Usage.Total.Known {
					t.Fatal("invented known usage")
				}
			}
		})
	}
}