internal/app/ai/service_test.go

Rama main · réplica de lectura del repositorio original.

package ai

import (
	"context"
	"errors"
	"strings"
	"testing"
	"time"

	"telesrv/internal/domain"
	"telesrv/internal/store/memory"
)

type fakeProvider struct {
	name string
	text string
	err  error
	seen ProviderRequest
}

func (p *fakeProvider) Name() string {
	if p.name == "" {
		return "fake"
	}
	return p.name
}

func (p *fakeProvider) Compose(_ context.Context, req ProviderRequest) (domain.AIComposeText, error) {
	p.seen = req
	if p.err != nil {
		return domain.AIComposeText{}, p.err
	}
	return domain.AIComposeText{Text: p.text}, nil
}

type fakeStreamingProvider struct {
	fakeProvider
	chunks []string
	final  string
}

func (p *fakeStreamingProvider) ComposeStream(_ context.Context, req ProviderRequest, emit func(domain.AIComposeText) error) (domain.AIComposeText, error) {
	p.seen = req
	if p.err != nil {
		return domain.AIComposeText{}, p.err
	}
	for _, chunk := range p.chunks {
		if emit != nil {
			if err := emit(domain.AIComposeText{Text: chunk}); err != nil {
				return domain.AIComposeText{}, err
			}
		}
	}
	final := p.final
	if final == "" && len(p.chunks) > 0 {
		final = p.chunks[len(p.chunks)-1]
	}
	return domain.AIComposeText{Text: final}, nil
}

type denyLimiter struct{}

func (denyLimiter) Allow(context.Context, string, int, time.Duration) (bool, int, error) {
	return false, 60, nil
}

func TestListTonesReturnsDefaultsAndHash(t *testing.T) {
	svc := NewService(memory.NewAIComposeStore())

	tones, notModified, err := svc.ListTones(context.Background(), 1001, 0)
	if err != nil || notModified {
		t.Fatalf("ListTones = notModified %v err %v", notModified, err)
	}
	if len(tones.Tones) == 0 {
		t.Fatal("ListTones returned no default tones; TDesktop would hide AI compose button")
	}
	if tones.Hash == 0 {
		t.Fatal("ListTones hash = 0, want stable non-zero hash")
	}
	_, notModified, err = svc.ListTones(context.Background(), 1001, tones.Hash)
	if err != nil || !notModified {
		t.Fatalf("ListTones(hash) = notModified %v err %v, want notModified", notModified, err)
	}
}

func TestDefaultTonePromptsDiscourageEcho(t *testing.T) {
	for _, tone := range DefaultTones() {
		if !strings.Contains(tone.Prompt, "Avoid returning the exact original text") {
			t.Fatalf("default tone %q prompt = %q, want echo guard", tone.Slug, tone.Prompt)
		}
	}
}

func TestDefaultTonesMatchOfficialDirectory(t *testing.T) {
	want := []struct {
		slug    string
		title   string
		emojiID int64
	}{
		{"formal", "Formal", defaultToneEmojiFormal},
		{"short", "Short", defaultToneEmojiShort},
		{"tribal", "Tribal", defaultToneEmojiTribal},
		{"corp", "Corp", defaultToneEmojiCorp},
		{"zen", "Zen", defaultToneEmojiZen},
		{"biblical", "Biblical", defaultToneEmojiBiblical},
		{"viking", "Viking", defaultToneEmojiViking},
	}
	got := DefaultTones()
	if len(got) != len(want) {
		t.Fatalf("DefaultTones len = %d, want %d", len(got), len(want))
	}
	for i, tone := range got {
		if tone.Slug != want[i].slug || tone.Title != want[i].title {
			t.Fatalf("DefaultTones[%d] = %s/%s, want %s/%s", i, tone.Slug, tone.Title, want[i].slug, want[i].title)
		}
		if !tone.Default {
			t.Fatalf("DefaultTones[%d].Default = false, want true", i)
		}
		if tone.EmojiID != want[i].emojiID {
			t.Fatalf("DefaultTones[%d].EmojiID = %d, want %d", i, tone.EmojiID, want[i].emojiID)
		}
		if tone.ExampleEnglish == nil {
			t.Fatalf("DefaultTones[%d].ExampleEnglish = nil", i)
		}
	}
}

func TestComposeCallsProviderWithInstruction(t *testing.T) {
	provider := &fakeProvider{text: "Please send the file when you have a moment."}
	svc := NewService(memory.NewAIComposeStore(), WithProvider(provider))

	got, err := svc.Compose(context.Background(), domain.AIComposeRequest{
		UserID: 1001,
		Text:   domain.AIComposeText{Text: "send file when free"},
		Tone:   domain.AIComposeToneRef{Kind: domain.AIComposeToneRefDefault, DefaultTone: "formal"},
	})
	if err != nil {
		t.Fatalf("Compose: %v", err)
	}
	if got.ResultText.Text != provider.text {
		t.Fatalf("Compose text = %q, want provider text", got.ResultText.Text)
	}
	if provider.seen.Instruction == "" || provider.seen.Tone.Slug != "formal" {
		t.Fatalf("provider request = %#v, want formal instruction", provider.seen)
	}
	for _, want := range []string{
		"Produce a visibly revised variant",
		"Make the selected style visible",
		"Avoid returning the exact original text",
	} {
		if !strings.Contains(provider.seen.Instruction, want) {
			t.Fatalf("instruction = %q, missing %q", provider.seen.Instruction, want)
		}
	}
}

func TestComposeInstructionDoesNotAnswerDraftQuestions(t *testing.T) {
	provider := &fakeProvider{text: "What is AI?"}
	svc := NewService(memory.NewAIComposeStore(), WithProvider(provider))

	if _, err := svc.Compose(context.Background(), domain.AIComposeRequest{
		UserID:    1001,
		Text:      domain.AIComposeText{Text: "what is AI"},
		Proofread: true,
	}); err != nil {
		t.Fatalf("Compose: %v", err)
	}
	for _, want := range []string{
		"not as a request, question, command, or chat message to answer",
		"Do not answer questions",
		"If the draft is a question, keep it as a question",
	} {
		if !strings.Contains(provider.seen.Instruction, want) {
			t.Fatalf("instruction = %q, missing %q", provider.seen.Instruction, want)
		}
	}
}

func TestComposeProofreadReturnsDiffText(t *testing.T) {
	provider := &fakeProvider{text: "Hello world."}
	svc := NewService(memory.NewAIComposeStore(), WithProvider(provider))

	got, err := svc.Compose(context.Background(), domain.AIComposeRequest{
		UserID:    1001,
		Text:      domain.AIComposeText{Text: "hello   world"},
		Proofread: true,
	})
	if err != nil {
		t.Fatalf("Compose: %v", err)
	}
	if got.DiffText == nil {
		t.Fatal("DiffText = nil, want proofread diff")
	}
	if got.DiffText.Text != "Hello world." || len(got.DiffText.Entities) != 1 {
		t.Fatalf("DiffText = %#v", got.DiffText)
	}
	ent := got.DiffText.Entities[0]
	if ent.Type != domain.MessageEntityDiffReplace || ent.Offset != 0 || ent.Length != 12 || ent.OldText != "hello   world" {
		t.Fatalf("diff entity = %#v", ent)
	}
}

func TestGetToneExampleUsesProviderForCustomTone(t *testing.T) {
	provider := &fakeProvider{text: "A crisp example."}
	store := memory.NewAIComposeStore()
	svc := NewService(store, WithProvider(provider))
	tone, err := svc.CreateTone(context.Background(), domain.AIComposeToneInput{
		UserID: 1001,
		Title:  "Crisp",
		Prompt: "Make it very crisp.",
	})
	if err != nil {
		t.Fatalf("CreateTone: %v", err)
	}
	got, err := svc.GetToneExample(context.Background(), 1001, domain.AIComposeToneRef{Kind: domain.AIComposeToneRefSlug, Slug: tone.Slug}, 2)
	if err != nil {
		t.Fatalf("GetToneExample: %v", err)
	}
	if got.To.Text != provider.text {
		t.Fatalf("example to = %q, want provider text", got.To.Text)
	}
	if provider.seen.Instruction == "" || provider.seen.Tone.ID != tone.ID {
		t.Fatalf("provider request = %#v, want custom tone instruction", provider.seen)
	}
}

func TestComposeRateLimited(t *testing.T) {
	svc := NewService(memory.NewAIComposeStore(), WithRateLimiter(denyLimiter{}, 1, time.Minute))
	_, err := svc.Compose(context.Background(), domain.AIComposeRequest{
		UserID:    1001,
		Text:      domain.AIComposeText{Text: "please polish this"},
		Proofread: true,
	})
	if !errors.Is(err, domain.ErrAIComposeRateLimited) {
		t.Fatalf("Compose err = %v, want ErrAIComposeRateLimited", err)
	}
}

func TestGenerateTextUsesProviderInstruction(t *testing.T) {
	provider := &fakeProvider{text: "Thanks for reaching out."}
	svc := NewService(memory.NewAIComposeStore(), WithProvider(provider))

	got, err := svc.GenerateText(context.Background(), domain.AITextGenerationRequest{
		UserID:      1001,
		Text:        domain.AIComposeText{Text: "hello"},
		Instruction: "Reply as the business owner.",
	})
	if err != nil {
		t.Fatalf("GenerateText: %v", err)
	}
	if got.Text != provider.text {
		t.Fatalf("GenerateText = %q, want provider text", got.Text)
	}
	if provider.seen.Instruction != "Reply as the business owner." || provider.seen.Request.Text.Text != "hello" {
		t.Fatalf("provider request = %#v", provider.seen)
	}
	if provider.seen.Purpose != ProviderPurposeTextGeneration {
		t.Fatalf("provider purpose = %q, want text generation", provider.seen.Purpose)
	}
}

func TestGenerateTextStreamUsesStreamingProvider(t *testing.T) {
	provider := &fakeStreamingProvider{
		chunks: []string{"Hel", "Hello"},
		final:  "Hello",
	}
	svc := NewService(memory.NewAIComposeStore(), WithProvider(provider))

	var chunks []string
	got, err := svc.GenerateTextStream(context.Background(), domain.AITextGenerationRequest{
		UserID:      1001,
		Text:        domain.AIComposeText{Text: "hello"},
		Instruction: "Reply as an assistant.",
	}, func(text domain.AIComposeText) error {
		chunks = append(chunks, text.Text)
		return nil
	})
	if err != nil {
		t.Fatalf("GenerateTextStream: %v", err)
	}
	if got.Text != "Hello" {
		t.Fatalf("final text = %q, want Hello", got.Text)
	}
	if len(chunks) != 2 || chunks[0] != "Hel" || chunks[1] != "Hello" {
		t.Fatalf("chunks = %#v", chunks)
	}
	if provider.seen.Instruction != "Reply as an assistant." || provider.seen.Request.Text.Text != "hello" {
		t.Fatalf("provider request = %#v", provider.seen)
	}
	if provider.seen.Purpose != ProviderPurposeTextGeneration {
		t.Fatalf("provider purpose = %q, want text generation", provider.seen.Purpose)
	}
}

func TestGenerateTextStreamFallsBackToNonStreamingProvider(t *testing.T) {
	provider := &fakeProvider{text: "One-shot answer."}
	svc := NewService(memory.NewAIComposeStore(), WithProvider(provider))

	var chunks []string
	got, err := svc.GenerateTextStream(context.Background(), domain.AITextGenerationRequest{
		UserID:      1001,
		Text:        domain.AIComposeText{Text: "hello"},
		Instruction: "Reply.",
	}, func(text domain.AIComposeText) error {
		chunks = append(chunks, text.Text)
		return nil
	})
	if err != nil {
		t.Fatalf("GenerateTextStream: %v", err)
	}
	if got.Text != "One-shot answer." || len(chunks) != 1 || chunks[0] != "One-shot answer." {
		t.Fatalf("final=%q chunks=%#v", got.Text, chunks)
	}
}

func TestGenerateTextStreamDoesNotFallbackToLocalEcho(t *testing.T) {
	provider := &fakeStreamingProvider{
		fakeProvider: fakeProvider{err: domain.ErrAIComposeProviderTimeout},
	}
	svc := NewService(memory.NewAIComposeStore(), WithProviders(provider, LocalProvider{}))

	var chunks []string
	_, err := svc.GenerateTextStream(context.Background(), domain.AITextGenerationRequest{
		UserID:      1001,
		Text:        domain.AIComposeText{Text: "User: secret prompt\nAssistant: hidden reply\nUser: hello"},
		Instruction: "Reply.",
	}, func(text domain.AIComposeText) error {
		chunks = append(chunks, text.Text)
		return nil
	})
	if !errors.Is(err, domain.ErrAIComposeProviderTimeout) {
		t.Fatalf("GenerateTextStream err = %v, want provider timeout", err)
	}
	if len(chunks) != 0 {
		t.Fatalf("chunks = %#v, want no local prompt echo", chunks)
	}
}

func TestGenerateTextRejectsLocalOnlyProvider(t *testing.T) {
	svc := NewService(memory.NewAIComposeStore(), WithProvider(LocalProvider{}))

	_, err := svc.GenerateText(context.Background(), domain.AITextGenerationRequest{
		UserID:      1001,
		Text:        domain.AIComposeText{Text: "User: hello"},
		Instruction: "Reply.",
	})
	if !errors.Is(err, domain.ErrAIComposeProviderUnavailable) {
		t.Fatalf("GenerateText err = %v, want provider unavailable", err)
	}
}

func TestCustomToneCRUDAndSave(t *testing.T) {
	store := memory.NewAIComposeStore()
	svc := NewService(store, WithClock(func() time.Time { return time.Unix(100, 0) }))

	tone, err := svc.CreateTone(context.Background(), domain.AIComposeToneInput{
		UserID:        1001,
		DisplayAuthor: true,
		Title:         "Sharp",
		Prompt:        "Make it direct and crisp.",
	})
	if err != nil {
		t.Fatalf("CreateTone: %v", err)
	}
	if tone.ID == 0 || tone.AccessHash == 0 || tone.Slug == "" || !tone.Creator || tone.AuthorID != 1001 {
		t.Fatalf("created tone = %#v", tone)
	}
	newTitle := "Brief"
	updated, err := svc.UpdateTone(context.Background(), domain.AIComposeToneUpdate{
		UserID: 1001,
		Ref: domain.AIComposeToneRef{
			Kind:       domain.AIComposeToneRefID,
			ID:         tone.ID,
			AccessHash: tone.AccessHash,
		},
		Title: &newTitle,
	})
	if err != nil {
		t.Fatalf("UpdateTone: %v", err)
	}
	if updated.Title != newTitle {
		t.Fatalf("updated title = %q, want %q", updated.Title, newTitle)
	}
	if err := svc.SaveTone(context.Background(), 2002, domain.AIComposeToneRef{Kind: domain.AIComposeToneRefSlug, Slug: tone.Slug}, false); err != nil {
		t.Fatalf("SaveTone: %v", err)
	}
	other, _, err := svc.ListTones(context.Background(), 2002, 0)
	if err != nil {
		t.Fatalf("ListTones other: %v", err)
	}
	var found bool
	for _, item := range other.Tones {
		if item.ID == tone.ID && item.Saved && !item.Creator {
			found = true
		}
	}
	if !found {
		t.Fatalf("saved tone not visible for other user: %#v", other.Tones)
	}
	if err := svc.DeleteTone(context.Background(), 1001, domain.AIComposeToneRef{Kind: domain.AIComposeToneRefID, ID: tone.ID, AccessHash: tone.AccessHash}); err != nil {
		t.Fatalf("DeleteTone: %v", err)
	}
	if _, err := svc.GetTone(context.Background(), 1001, domain.AIComposeToneRef{Kind: domain.AIComposeToneRefSlug, Slug: tone.Slug}); !errors.Is(err, domain.ErrAIComposeToneNotFound) {
		t.Fatalf("GetTone after delete err = %v, want ErrAIComposeToneNotFound", err)
	}
}