bub/tests/test_tape_service.py at 76aefeeef79fad51a376252550e20b1e02351e6e · bubbuild/bub · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
from dataclasses import dataclass

from bub.tape.service import TapeService


@dataclass
class FakeEntry:
    id: int
    kind: str
    payload: dict[str, object]
    meta: dict[str, object]


class FakeTape:
    class _Query:
        def __init__(self, tape: "FakeTape") -> None:
            self._tape = tape

        def all(self) -> list[FakeEntry]:
            return list(self._tape.entries)

    def __init__(self) -> None:
        self.name = "fake"
        self.entries: list[FakeEntry] = [
            FakeEntry(
                id=1,
                kind="anchor",
                payload={"name": "session/start", "state": {"owner": "human"}},
                meta={},
            )
        ]
        self.reset_calls = 0
        self.query = self._Query(self)

    def handoff(self, name: str, state: dict[str, object] | None = None) -> list[FakeEntry]:
        entry = FakeEntry(
            id=len(self.entries) + 1,
            kind="anchor",
            payload={"name": name, "state": state or {}},
            meta={},
        )
        self.entries.append(entry)
        return [entry]

    def reset(self) -> None:
        self.reset_calls += 1
        self.entries = []


def test_reset_rebuilds_bootstrap_anchor() -> None:
    service = TapeService.__new__(TapeService)
    fake_tape = FakeTape()
    service._tape = fake_tape  # type: ignore[attr-defined]
    service._store = None  # type: ignore[attr-defined]

    result = service.reset()

    assert result == "ok"
    assert fake_tape.reset_calls == 1
    anchors = [entry for entry in fake_tape.entries if entry.kind == "anchor"]
    assert len(anchors) == 1
    assert anchors[0].payload["name"] == "session/start"


def test_search_supports_fuzzy_typo_matching() -> None:
    service = TapeService.__new__(TapeService)
    fake_tape = FakeTape()
    fake_tape.entries.extend((
        FakeEntry(
            id=2,
            kind="message",
            payload={"role": "assistant", "content": "Please review the database migration plan."},
            meta={"source": "assistant"},
        ),
        FakeEntry(
            id=3,
            kind="message",
            payload={"role": "assistant", "content": "Unrelated note"},
            meta={},
        ),
    ))
    service._tape = fake_tape  # type: ignore[attr-defined]

    matches = service.search("databse migrtion", limit=5)

    assert len(matches) == 1
    assert matches[0].id == 2


def test_search_respects_limit_for_exact_match() -> None:
    service = TapeService.__new__(TapeService)
    fake_tape = FakeTape()
    fake_tape.entries.extend((
        FakeEntry(
            id=2,
            kind="message",
            payload={"role": "assistant", "content": "Alpha report generated"},
            meta={},
        ),
        FakeEntry(
            id=3,
            kind="message",
            payload={"role": "assistant", "content": "Alpha follow-up details"},
            meta={},
        ),
    ))
    service._tape = fake_tape  # type: ignore[attr-defined]

    matches = service.search("alpha", limit=1)

    assert len(matches) == 1
    assert matches[0].id == 3