Agent-System-Backend/internal/execenv/broker_test.go

package execenv

import (
	"archive/zip"
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"os"
	"path/filepath"
	"testing"
)

type memStore struct {
	reqs map[string]Request
	bind map[string]Binding
}

func newMem() *memStore { return &memStore{reqs: map[string]Request{}, bind: map[string]Binding{}} }

func (m *memStore) InsertExecEnvRequest(r Request) error {
	m.reqs[r.ID] = r
	return nil
}
func (m *memStore) GetExecEnvRequest(id string) (Request, error) { return m.reqs[id], nil }
func (m *memStore) ListExecEnvRequests(conv string, limit int) ([]Request, error) {
	var out []Request
	for _, r := range m.reqs {
		if conv == "" || r.ConversationID == conv {
			out = append(out, r)
		}
	}
	return out, nil
}
func (m *memStore) UpdateExecEnvRequest(id, status, layerID, detail, approver string) error {
	r := m.reqs[id]
	r.Status, r.LayerID, r.Detail, r.Approver = status, layerID, detail, approver
	m.reqs[id] = r
	return nil
}
func (m *memStore) GetExecEnvBinding(conv string) (Binding, error) { return m.bind[conv], nil }
func (m *memStore) PutExecEnvBinding(b Binding) error {
	m.bind[b.ConversationID] = b
	return nil
}
func (m *memStore) ListExecEnvBindings() ([]Binding, error) {
	var out []Binding
	for _, b := range m.bind {
		out = append(out, b)
	}
	return out, nil
}
func (m *memStore) InsertExecEnvAudit(string, string, string, string, string, string) error {
	return nil
}

type scripted struct{ byReq map[string]Result }

func (s scripted) Resolve(_ context.Context, reqs []string, _ Lock) (Result, error) {
	if len(reqs) == 0 {
		return Result{}, nil
	}
	return s.byReq[reqs[0]], nil
}
func (s scripted) Fetch(context.Context, string, LockedPkg) (string, error) {
	return "", os.ErrNotExist
}

func TestConversationsGetDistinctLayers(t *testing.T) {
	root, err := Open(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	cat := DefaultCatalog()
	cat.Allow = append(cat.Allow, CatalogEntry{Name: "foo", Reason: "test"})
	raw, _ := json.Marshal(cat)
	if err := os.WriteFile(root.CatalogPath(), raw, 0o644); err != nil {
		t.Fatal(err)
	}
	baseSum := placeWheel(t, root, "bar.py")
	baseID, err := root.BuildLayer(Lock{
		Schema: 1, Kind: "base", PythonTag: "cp314", Platform: "any",
		Packages: []LockedPkg{{Name: "bar", Version: "1", Filename: baseSum + ".whl", SHA256: baseSum}},
	})
	if err != nil {
		t.Fatal(err)
	}
	if err := root.SwitchCurrent(baseID); err != nil {
		t.Fatal(err)
	}
	sum1 := placeWheel(t, root, "foo1.py")
	sum2 := placeWheel(t, root, "foo2.py")
	broker := NewBroker(root, newMem(), scripted{byReq: map[string]Result{
		"foo==1": {Packages: []LockedPkg{{Name: "foo", Version: "1", Filename: sum1 + ".whl", SHA256: sum1}}},
		"foo==2": {Packages: []LockedPkg{{Name: "foo", Version: "2", Filename: sum2 + ".whl", SHA256: sum2}}},
	}})
	a, err := broker.RequestPackages(context.Background(), "conv-a", "owner", "need 1", []string{"foo==1"}, "")
	if err != nil || a.Status != "ready" {
		t.Fatalf("A: %+v %v", a, err)
	}
	b, err := broker.RequestPackages(context.Background(), "conv-b", "owner", "need 2", []string{"foo==2"}, "")
	if err != nil || b.Status != "ready" {
		t.Fatalf("B: %+v %v", b, err)
	}
	if a.LayerID == b.LayerID {
		t.Fatalf("layers should differ: %s", a.LayerID)
	}
	cur, err := root.CurrentBase()
	if err != nil || cur != baseID {
		t.Fatalf("base changed: %s -> %s (%v)", baseID, cur, err)
	}
	extraA, isoA := broker.LayersFor("conv-a")
	extraB, isoB := broker.LayersFor("conv-b")
	if isoA != "" || isoB != "" || len(extraA) != 1 || len(extraB) != 1 || extraA[0] == extraB[0] {
		t.Fatalf("mounts A=%v %s B=%v %s", extraA, isoA, extraB, isoB)
	}
}

func TestRequestGrades(t *testing.T) {
	root, err := Open(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	broker := NewBroker(root, newMem(), scripted{byReq: map[string]Result{
		"onlysrc": {SDist: []string{"onlysrc"}},
		"secret":  {Packages: []LockedPkg{{Name: "secret", Version: "1", SHA256: "abc"}}},
	}})
	sdist, err := broker.RequestPackages(context.Background(), "c", "o", "", []string{"onlysrc"}, "")
	if err != nil || sdist.Status != "rejected" {
		t.Fatalf("sdist: %+v %v", sdist, err)
	}
	outside, err := broker.RequestPackages(context.Background(), "c", "o", "", []string{"secret"}, "")
	if err != nil || outside.Status != "pending_approval" {
		t.Fatalf("allowlist: %+v %v", outside, err)
	}
}

func placeWheel(t *testing.T, root *Root, member string) string {
	t.Helper()
	buf := &bytes.Buffer{}
	zw := zip.NewWriter(buf)
	f, err := zw.Create(member)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := f.Write([]byte("x = 1\n")); err != nil {
		t.Fatal(err)
	}
	if err := zw.Close(); err != nil {
		t.Fatal(err)
	}
	sum := sha256.Sum256(buf.Bytes())
	hexSum := hex.EncodeToString(sum[:])
	if err := os.WriteFile(filepath.Join(root.Wheelhouse(), hexSum+".whl"), buf.Bytes(), 0o644); err != nil {
		t.Fatal(err)
	}
	return hexSum
}

func osWriteWheel(path, member string) error {
	buf := &bytes.Buffer{}
	zw := zip.NewWriter(buf)
	if _, err := zw.Create(member); err != nil {
		return err
	}
	if err := zw.Close(); err != nil {
		return err
	}
	return os.WriteFile(path, buf.Bytes(), 0o644)
}