xinmi/internal/workbench/memory_snapshot_test.go

package workbench

import (
	"context"
	"strings"
	"testing"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
)

type snapshotMemoryStub struct {
	interfaces.MemoryService
	calls int
}

func (m *snapshotMemoryStub) Recall(context.Context, string) interfaces.MemoryRecall {
	m.calls++
	return interfaces.MemoryRecall{Prompt: "当前用户的私有记忆"}
}

func TestFrozenAgentCannotReenableDisabledMemory(t *testing.T) {
	for _, state := range []string{"live_disabled", "snapshot_disabled", "both_enabled"} {
		t.Run(state, func(t *testing.T) {
			s, ctx, _, _ := generationHarness(t)
			disabled, enabled := false, true
			live := &enabled
			saved := &enabled
			if state == "live_disabled" {
				live = &disabled
			}
			if state == "snapshot_disabled" {
				saved = &disabled
			}
			s.agents = generationAgentStub{config: types.CustomAgentConfig{MemoryEnabled: live}}
			memories := &snapshotMemoryStub{}
			s.memories = memories
			ctx = context.WithValue(ctx, agentSnapshotsKey{}, map[string]AgentSnapshot{
				"author": {Name: "Author", ModelID: "model", MemoryEnabled: saved, OutputLimit: 2000},
			})
			completion, err := s.prepareCompletion(ctx, "author", "写作", "编写本段")
			if err != nil {
				t.Fatal(err)
			}
			wantRecall := state == "both_enabled"
			if (memories.calls == 1) != wantRecall || strings.Contains(completion.messages[0].Content, "当前用户的私有记忆") != wantRecall {
				t.Fatalf("memory preference ignored: state=%s calls=%d", state, memories.calls)
			}
		})
	}
}