HomeAgent/internal/memory/pipeline/splitter_test.go

package pipeline

import (
	"context"
	"errors"
	"path/filepath"
	"testing"
	"time"

	"github.com/JianFeeeee/HomeAgent/internal/memory"
	"github.com/JianFeeeee/HomeAgent/internal/memory/distill"
)

// newTestDistiller 建一个带真实图库的 Distiller（nil db 的那种测不了入库）。
func newTestDistiller(t *testing.T) (*Distiller, func()) {
	t.Helper()
	db, err := memory.NewGraphDB(filepath.Join(t.TempDir(), "test.db"))
	if err != nil {
		t.Fatalf("NewGraphDB: %v", err)
	}
	d := NewDistiller(db, t.TempDir(), DistillerConfig{Interval: time.Hour, BatchSize: 50})
	d.ctx, d.cancel = context.WithCancel(context.Background())
	return d, func() {
		d.cancel()
		db.Close()
	}
}

func nowForTest() time.Time { return time.Now() }

// relationCount 读活动关系数。
func relationCount(t *testing.T, d *Distiller) int {
	t.Helper()
	data, err := d.db.Introspect()
	if err != nil {
		t.Fatalf("Introspect: %v", err)
	}
	if n, ok := data["relation_count"].(int); ok {
		return n
	}
	return -1
}

// stubSplitter 记录调用参数，返回预置结果。
type stubSplitter struct {
	called  int
	gotText string
	triples []memory.Triple
	err     error
	callErr error
}

func (s *stubSplitter) Split(_ context.Context, record string) ([]memory.Triple, error) {
	s.called++
	s.gotText = record
	if s.callErr != nil {
		return nil, s.callErr
	}
	return s.triples, s.err
}

// ★ splitter 注入后应被真正使用（否则接了等于没接）。
func TestDistillBatch_优先用Splitter(t *testing.T) {
	d, cleanup := newTestDistiller(t)
	defer cleanup()

	want := []memory.Triple{{Subject: "第183批", Relation: "停机时长", Object: "7分"}}
	stub := &stubSplitter{triples: want}
	d.SetSplitter(stub)

	ok := d.distillBatch([]RawRecord{{ID: 1, SessionID: "sess", Role: "user",
		Content: "第183批周三凌晨1点·停机7分", CreatedAt: nowForTest()}})
	if !ok {
		t.Fatal("distillBatch 应成功")
	}
	if stub.called != 1 {
		t.Fatalf("splitter 应被调用 1 次，实际 %d", stub.called)
	}
	if !contains(stub.gotText, "停机7分") {
		t.Errorf("splitter 未收到记录内容：%q", stub.gotText)
	}
	// 三元组应真的进了图
	if n := relationCount(t, d); n == 0 {
		t.Error("splitter 产出的三元组未入库")
	}
}

// splitter 返回空但成功 = 「这条记录没有可拆字段」，应视为成功
// （蒸馏标记完成，不重试）。
func TestDistillBatch_splitter空结果视为成功(t *testing.T) {
	d, cleanup := newTestDistiller(t)
	defer cleanup()
	stub := &stubSplitter{triples: nil}
	d.SetSplitter(stub)

	ok := d.distillBatch([]RawRecord{{ID: 1, SessionID: "s", Role: "user",
		Content: "第129~143批均仅评审通过"}})
	if !ok {
		t.Fatal("splitter 返回 0 条应视为成功（无字段可拆是正常结果）")
	}
	if stub.called != 1 {
		t.Fatalf("splitter 应被调用，实际 %d", stub.called)
	}
}

// ★ splitter 报错 → 回退 jieba 路，且 distillBatch 仍成功（不停摆）。
func TestDistillBatch_splitter失败回退(t *testing.T) {
	d, cleanup := newTestDistiller(t)
	defer cleanup()
	stub := &stubSplitter{err: errors.New("model unavailable")}
	d.SetSplitter(stub)

	ok := d.distillBatch([]RawRecord{{ID: 1, SessionID: "s", Role: "user",
		Content: "第183批周三凌晨1点·停机7分"}})
	if !ok {
		t.Fatal("splitter 失败后应回退而非整体失败")
	}
	if stub.called != 1 {
		t.Fatalf("splitter 应被尝试过 1 次，实际 %d", stub.called)
	}
}

// 未注入 splitter 时行为不变（向后兼容）。
func TestDistillBatch_无splitter走原路径(t *testing.T) {
	d, cleanup := newTestDistiller(t)
	defer cleanup()
	d.SetSplitter(nil)

	ok := d.distillBatch([]RawRecord{{ID: 1, SessionID: "s", Role: "user",
		Content: "第183批周三凌晨1点·停机7分"}})
	if !ok {
		t.Fatal("未注入 splitter 时应仍能成功（走既有 jieba 路）")
	}
}

// ★ 变异自证：若 extractTriples 忽略 splitter 恒走 jieba，
// 「优先用 splitter」这条测试必须变红。
func TestExtractTriples_变异_忽略splitter则失败(t *testing.T) {
	d, cleanup := newTestDistiller(t)
	defer cleanup()
	stub := &stubSplitter{triples: []memory.Triple{{Subject: "x"}}}
	d.SetSplitter(stub)

	got := d.extractTriples("第183批 停机7分", "")
	if len(got) != 1 || got[0].Subject != "x" {
		t.Fatalf("extractTriples 应返回 splitter 的结果，实际 %+v", got)
	}
}

func contains(s, sub string) bool {
	return len(sub) == 0 || (len(s) >= len(sub) && indexOf(s, sub) >= 0)
}

func indexOf(s, sub string) int {
	for i := 0; i+len(sub) <= len(s); i++ {
		if s[i:i+len(sub)] == sub {
			return i
		}
	}
	return -1
}

// blockSplitter 实现 BlockSplitter（多了 Blocks 方法）。
type blockSplitter struct {
	stubSplitter
	payload  *distill.BlockPayload
	blockErr error
	calls    int
}

func (b *blockSplitter) Blocks(_ context.Context, record string) (*distill.BlockPayload, error) {
	b.calls++
	if b.blockErr != nil {
		return nil, b.blockErr
	}
	if b.payload != nil {
		return b.payload, nil
	}
	return &distill.BlockPayload{
		Sentence: record,
		Fields: []distill.FieldBlock{
			{Dimension: "停机时长", Value: "4分"},
		},
	}, nil
}

// ★ 块路径优先：实现 BlockSplitter 时，产出应落成块（不是 entities）。
func TestDistillBatch_块路径优先(t *testing.T) {
	d, cleanup := newTestDistiller(t)
	defer cleanup()
	bs := &blockSplitter{}
	d.SetSplitter(bs)

	ok := d.distillBatch([]RawRecord{{ID: 1, SessionID: "s", Role: "user",
		Content: "第112批 停机4分"}})
	if !ok {
		t.Fatal("distillBatch 应成功")
	}
	if bs.calls == 0 {
		t.Fatal("应走块路径")
	}
	blocks, err := d.db.MemoryBlocks()
	if err != nil {
		t.Fatal(err)
	}
	if len(blocks) == 0 {
		t.Fatal("块路径应写入块")
	}
	// ★ 关键：块路径**不写旧表 entities**。
	//
	// ★ 判据口径变了（2026-10-04）：原来断言「Recall 返回 0 个实体」。
	//   那是在 Recall 仍读旧表时的写法 —— Recall 切块之后它必然返回块
	//   （那正是块路径该做的事），于是判据恒红。
	//
	//   判据的**意图**是「没有走旧表双写」，所以直接查旧表本身：
	//   与实现无关，且比间接断言可靠。
	legacy, err := d.db.LegacyEntities(0, 1)
	if err != nil {
		t.Fatalf("LegacyEntities: %v", err)
	}
	if len(legacy) != 0 {
		t.Errorf("块路径不该写 entities（旧表 %d 个）: %+v", len(legacy), legacy)
	}
}

// ★ 块路失败时落回 Triple 路，而不是让整批重试。
func TestDistillBatch_块路失败落回Triple(t *testing.T) {
	d, cleanup := newTestDistiller(t)
	defer cleanup()
	bs := &blockSplitter{
		stubSplitter: stubSplitter{triples: []memory.Triple{{Subject: "s", Relation: "r", Object: "o"}}},
		blockErr:     errors.New("model timeout"),
	}
	d.SetSplitter(bs)

	ok := d.distillBatch([]RawRecord{{ID: 1, SessionID: "s", Role: "user",
		Content: "第112批 停机4分"}})
	if !ok {
		t.Fatal("块路失败应落回 Triple 路并成功，而非整批失败（那会无限重试）")
	}
}

// 不带 BlockSplitter 能力（只有 Split）时走旧路，不 panic。
func TestDistillBatch_仅Split能力走旧路(t *testing.T) {
	d, cleanup := newTestDistiller(t)
	defer cleanup()
	st := &stubSplitter{triples: []memory.Triple{{Subject: "s", Relation: "r", Object: "o"}}}
	d.SetSplitter(st)

	if ok := d.distillBatch([]RawRecord{{ID: 1, SessionID: "s", Role: "user",
		Content: "第112批 停机4分"}}); !ok {
		t.Fatal("应成功")
	}
	if st.called == 0 {
		t.Fatal("应走 Triple 路")
	}
}