email_agent/internal/reply/draft_test.go

package reply

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

	"gitea.lab.audasmedia.com.au/sam/email_agent/internal/store"
)

func TestTemplateForAction(t *testing.T) {
	if got := TemplateForAction("reply").ID; got != "answer" {
		t.Errorf("TemplateForAction(reply) = %q, want answer", got)
	}
	if got := TemplateForAction("keep").ID; got != "quick-ack" {
		t.Errorf("TemplateForAction(keep) = %q, want quick-ack", got)
	}
	if _, ok := TemplateByID("DECLINE"); !ok {
		t.Error("TemplateByID should match case-insensitively")
	}
	if _, ok := TemplateByID("nonsense"); ok {
		t.Error("TemplateByID should reject an unknown template")
	}
}

func TestBuildPromptIncludesContext(t *testing.T) {
	req := Request{
		Message: store.Message{
			FromAddr:      "customer@example.test",
			SourceAddress: "hello@where-woof.com",
			Subject:       "Booking for Saturday",
			BodyText:      "Can we book Saturday at 10?",
		},
		Category:    "Action",
		Importance:  "high",
		Action:      "reply",
		Instruction: "Confirm 10am works",
		Signer:      "Sam",
	}
	tmpl, _ := TemplateByID("confirm")
	system, user := BuildPrompt(req, tmpl)

	if !strings.Contains(system, "Return only the reply body") {
		t.Errorf("system prompt missing the output constraint: %q", system)
	}
	for _, want := range []string{
		"Confirm an arrangement",
		"Confirm 10am works",
		"category=Action",
		"importance=high",
		"suggested_action=reply",
		"customer@example.test",
		"hello@where-woof.com",
		"Booking for Saturday",
		"Can we book Saturday at 10?",
	} {
		if !strings.Contains(user, want) {
			t.Errorf("user prompt missing %q\n%s", want, user)
		}
	}
}

func TestCleanDraftStripsFenceAndSubject(t *testing.T) {
	if got := cleanDraft("```\nHello there.\n```", 1000); got != "Hello there." {
		t.Errorf("cleanDraft fence = %q", got)
	}
	if got := cleanDraft("Subject: Re: hi\n\nHello there.", 1000); got != "Hello there." {
		t.Errorf("cleanDraft subject = %q", got)
	}
}

func TestGeneratorDraftUsesTemplate(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewEncoder(w).Encode(chatResponse{
			Choices: []struct {
				Message chatMessage `json:"message"`
			}{{Message: chatMessage{Role: "assistant", Content: "Sure, Saturday works.\n\nSam"}}},
		})
	}))
	defer srv.Close()

	g := NewGenerator(NewClient(srv.URL, "key", "auto/best-chat", slog.New(slog.DiscardHandler)),
		slog.New(slog.DiscardHandler))
	body, templateID, err := g.Draft(context.Background(), Request{
		Message: store.Message{FromAddr: "a@b.test", Subject: "hi", BodyText: "?"},
		Action:  "reply",
	})
	if err != nil {
		t.Fatalf("Draft() error = %v", err)
	}
	if templateID != "answer" {
		t.Errorf("template = %q, want answer (the default for action=reply)", templateID)
	}
	if !strings.Contains(body, "Saturday") {
		t.Errorf("body = %q", body)
	}
}

func TestGeneratorDraftEmptyReplyIsAnError(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewEncoder(w).Encode(chatResponse{
			Choices: []struct {
				Message chatMessage `json:"message"`
			}{{Message: chatMessage{Role: "assistant", Content: "   "}}},
		})
	}))
	defer srv.Close()

	g := NewGenerator(NewClient(srv.URL, "key", "m", slog.New(slog.DiscardHandler)),
		slog.New(slog.DiscardHandler))
	if _, _, err := g.Draft(context.Background(), Request{Message: store.Message{Subject: "x"}}); err == nil {
		t.Fatal("Draft() error = nil, want an error for an empty reply")
	}
}

func TestBuildPromptIncludesRules(t *testing.T) {
	req := Request{
		Message: store.Message{FromAddr: "customer@example.test", Subject: "hi", BodyText: "?"},
		Action:  "reply",
		Rules:   `from contains "example.test" => category=Action`,
	}
	tmpl, _ := TemplateByID("answer")
	_, user := BuildPrompt(req, tmpl)
	if !strings.Contains(user, "USER RULES") || !strings.Contains(user, `from contains "example.test"`) {
		t.Errorf("user prompt does not include the rules document:\n%s", user)
	}
}