Demonstrate overlay model routing in prototype

examples/library-adoption/workflow_test.go

@@ -1,9 +1,11 @@
 package main
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
+	"io"
 	"net/http"
 	"strings"
 	"sync"
@@ -13,6 +15,74 @@ import (
 	pk "gitea.maximumdirect.net/eric/promptkit"
 )
 
+func TestDemonstrationSelectsOutboundModels(t *testing.T) {
+	for _, tc := range []struct {
+		name                          string
+		args                          []string
+		overlay, proposal, assessment string
+	}{
+		{"base", nil, "", "synthetic-grammar", "synthetic-review"},
+		{"testing", []string{"--overlay", "testing"}, "testing", "synthetic-grammar-testing", "synthetic-review-testing"},
+		{"proposal override", []string{"--overlay", "testing", "--proposal-profile", "proposal"}, "testing", "synthetic-grammar", "synthetic-review-testing"},
+		{"assessment override", []string{"--overlay", "testing", "--validation-profile", "assessment"}, "testing", "synthetic-grammar-testing", "synthetic-review"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Setenv("AUDITA_PROTOTYPE_KEY", "synthetic-test-secret")
+			var mu sync.Mutex
+			var models []string
+			base := demonstrationClient()
+			client := &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
+				body, err := io.ReadAll(req.Body)
+				if err != nil {
+					return nil, err
+				}
+				var wire struct {
+					Model string `json:"model"`
+				}
+				if err := json.Unmarshal(body, &wire); err != nil {
+					return nil, err
+				}
+				mu.Lock()
+				models = append(models, wire.Model)
+				mu.Unlock()
+				req.Body = io.NopCloser(bytes.NewReader(body))
+				return base.Transport.RoundTrip(req)
+			})}
+			var out bytes.Buffer
+			if err := run(tc.args, &out, client); err != nil {
+				t.Fatal(err)
+			}
+			var got struct {
+				Overlay    string                                   `json:"overlay"`
+				Proposal   struct{ Profile, Backend, Model string } `json:"proposal"`
+				Assessment struct{ Profile, Backend, Model string } `json:"assessment"`
+				Result     WorkflowResult                           `json:"result"`
+				Attempts   uint64                                   `json:"attempts"`
+				Retries    uint64                                   `json:"transient_retries"`
+				Repairs    uint64                                   `json:"structural_repairs"`
+			}
+			if err := json.Unmarshal(out.Bytes(), &got); err != nil {
+				t.Fatal(err)
+			}
+			if got.Overlay != tc.overlay || got.Proposal.Model != tc.proposal || got.Assessment.Model != tc.assessment || got.Proposal.Backend != "offline" || got.Assessment.Backend != "offline" {
+				t.Fatalf("selection = %+v", got)
+			}
+			profiles := map[string]string{"synthetic-grammar": "proposal", "synthetic-review": "assessment", "synthetic-grammar-testing": "proposal-testing", "synthetic-review-testing": "assessment-testing"}
+			if got.Proposal.Profile != profiles[tc.proposal] || got.Assessment.Profile != profiles[tc.assessment] {
+				t.Fatalf("reported profiles = %q, %q", got.Proposal.Profile, got.Assessment.Profile)
+			}
+			if got.Result.Accepted != 1 || got.Attempts != 4 || got.Retries != 0 || got.Repairs != 2 {
+				t.Fatalf("outcome = %+v", got)
+			}
+			mu.Lock()
+			defer mu.Unlock()
+			if len(models) != 4 || models[0] != tc.proposal || models[1] != tc.proposal || models[2] != tc.assessment || models[3] != tc.assessment {
+				t.Fatalf("outbound models = %v", models)
+			}
+		})
+	}
+}
+
 type step struct {
 	status  int
 	content string