solpbc.org
personal memory agent
1# SPDX-License-Identifier: AGPL-3.0-only
2# Copyright (c) 2026 sol pbc
3
4"""Tests for observe/extract.py frame selection logic."""
5
6from unittest.mock import patch
7
8from observe.extract import (
9 DEFAULT_MAX_EXTRACTIONS,
10 _build_extraction_guidance,
11 _fallback_select_frames,
12 select_frames_for_extraction,
13)
14
15
16def _make_frames(count: int, start_id: int = 1) -> list[dict]:
17 """Create test frame data."""
18 return [
19 {"frame_id": i, "timestamp": float(i), "analysis": {"primary": "code"}}
20 for i in range(start_id, start_id + count)
21 ]
22
23
24def test_default_max_extractions():
25 """Test default max extractions value."""
26 assert DEFAULT_MAX_EXTRACTIONS == 20
27
28
29def test_empty_list():
30 """Test empty input returns empty output."""
31 result = select_frames_for_extraction([])
32 assert result == []
33
34
35def test_single_frame():
36 """Test single frame is always selected."""
37 frames = _make_frames(1)
38 result = select_frames_for_extraction(frames, max_extractions=5)
39 assert result == [1]
40
41
42def test_fewer_than_max_returns_all():
43 """Test that frames under the limit are all returned."""
44 frames = _make_frames(5)
45 result = select_frames_for_extraction(frames, max_extractions=10)
46 assert result == [1, 2, 3, 4, 5]
47
48
49def test_exactly_max_returns_all():
50 """Test that exactly max frames returns all."""
51 frames = _make_frames(10)
52 result = select_frames_for_extraction(frames, max_extractions=10)
53 assert result == list(range(1, 11))
54
55
56def test_more_than_max_returns_around_max():
57 """Test that more than max frames returns approximately max count."""
58 frames = _make_frames(30)
59 result = select_frames_for_extraction(frames, max_extractions=5)
60 # May be max or max+1 if first frame wasn't in random selection
61 assert 5 <= len(result) <= 6
62
63
64def test_first_frame_always_included():
65 """Test that first frame is always in selection."""
66 frames = _make_frames(100)
67 # Run multiple times to account for randomness
68 for _ in range(10):
69 result = select_frames_for_extraction(frames, max_extractions=10)
70 assert 1 in result, "First frame must always be included"
71
72
73def test_results_sorted():
74 """Test that results are sorted by frame_id."""
75 frames = _make_frames(50)
76 result = select_frames_for_extraction(frames, max_extractions=10)
77 assert result == sorted(result)
78
79
80def test_fallback_deterministic_for_small_lists():
81 """Test fallback is deterministic for lists under max."""
82 frames = _make_frames(5)
83 result1 = _fallback_select_frames(frames, max_extractions=10)
84 result2 = _fallback_select_frames(frames, max_extractions=10)
85 assert result1 == result2 == [1, 2, 3, 4, 5]
86
87
88def test_max_extractions_of_one():
89 """Test edge case of max_extractions=1."""
90 frames = _make_frames(10)
91 result = select_frames_for_extraction(frames, max_extractions=1)
92 # First frame always included, plus possibly one random
93 assert 1 in result
94 assert 1 <= len(result) <= 2
95
96
97def test_non_sequential_frame_ids():
98 """Test with non-sequential frame IDs."""
99 frames = [
100 {"frame_id": 5, "timestamp": 1.0, "analysis": {}},
101 {"frame_id": 10, "timestamp": 2.0, "analysis": {}},
102 {"frame_id": 15, "timestamp": 3.0, "analysis": {}},
103 ]
104 result = select_frames_for_extraction(frames, max_extractions=10)
105 assert result == [5, 10, 15]
106
107
108# --- Tests for _build_extraction_guidance ---
109
110
111def test_build_extraction_guidance_empty():
112 """Test extraction guidance with empty categories."""
113 result = _build_extraction_guidance({})
114 assert result == "No category-specific rules."
115
116
117def test_build_extraction_guidance_no_extraction_fields():
118 """Test extraction guidance when no categories have extraction field."""
119 categories = {
120 "code": {"description": "Code editors"},
121 "browsing": {"description": "Web browsing"},
122 }
123 result = _build_extraction_guidance(categories)
124 assert result == "No category-specific rules."
125
126
127def test_build_extraction_guidance_with_extraction():
128 """Test extraction guidance with valid extraction fields."""
129 categories = {
130 "code": {
131 "description": "Code editors",
132 "extraction": "Extract on file changes",
133 },
134 "browsing": {
135 "description": "Web browsing",
136 "extraction": "Extract on site changes",
137 },
138 "meeting": {"description": "Video calls"}, # No extraction field
139 }
140 result = _build_extraction_guidance(categories)
141 assert "- browsing: Extract on site changes" in result
142 assert "- code: Extract on file changes" in result
143 assert "meeting" not in result
144
145
146def test_build_extraction_guidance_sorted():
147 """Test extraction guidance is sorted alphabetically."""
148 categories = {
149 "zebra": {"extraction": "Z rule"},
150 "apple": {"extraction": "A rule"},
151 "mango": {"extraction": "M rule"},
152 }
153 result = _build_extraction_guidance(categories)
154 lines = result.split("\n")
155 assert lines[0].startswith("- apple:")
156 assert lines[1].startswith("- mango:")
157 assert lines[2].startswith("- zebra:")
158
159
160# --- Tests for AI selection with mocked generate ---
161
162
163def test_ai_selection_with_categories():
164 """Test that AI selection is used when categories are provided."""
165 frames = _make_frames(10)
166 categories = {"code": {"description": "Code editors"}}
167
168 # Mock generate to return specific frame IDs
169 with patch("think.models.generate") as mock_generate:
170 mock_generate.return_value = "[1, 3, 5, 7]"
171 result = select_frames_for_extraction(
172 frames, max_extractions=5, categories=categories
173 )
174
175 assert result == [1, 3, 5, 7]
176 mock_generate.assert_called_once()
177
178
179def test_ai_selection_filters_invalid_ids():
180 """Test that AI selection filters out invalid frame IDs."""
181 frames = _make_frames(5) # IDs 1-5
182 categories = {"code": {"description": "Code editors"}}
183
184 # Mock generate to return some invalid IDs
185 with patch("think.models.generate") as mock_generate:
186 mock_generate.return_value = "[1, 3, 99, 100, 5]" # 99, 100 are invalid
187 result = select_frames_for_extraction(
188 frames, max_extractions=10, categories=categories
189 )
190
191 assert result == [1, 3, 5]
192 assert 99 not in result
193 assert 100 not in result
194
195
196def test_ai_selection_hard_cap():
197 """Test that AI selection caps at 2x max_extractions."""
198 frames = _make_frames(50)
199 categories = {"code": {"description": "Code editors"}}
200
201 # Mock generate to return way too many IDs
202 many_ids = list(range(1, 51)) # 50 IDs
203 with patch("think.models.generate") as mock_generate:
204 mock_generate.return_value = str(many_ids)
205 result = select_frames_for_extraction(
206 frames, max_extractions=5, categories=categories
207 )
208
209 # Hard cap is 2 * 5 = 10, plus first frame guarantee
210 assert len(result) <= 11
211
212
213def test_ai_selection_fallback_on_error():
214 """Test that fallback is used when AI selection fails."""
215 frames = _make_frames(10)
216 categories = {"code": {"description": "Code editors"}}
217
218 # Mock generate to raise an exception
219 with patch("think.models.generate") as mock_generate:
220 mock_generate.side_effect = Exception("API error")
221 result = select_frames_for_extraction(
222 frames, max_extractions=5, categories=categories
223 )
224
225 # Should still get a valid result from fallback
226 assert len(result) >= 5
227 assert 1 in result # First frame always included
228
229
230def test_ai_selection_fallback_on_invalid_json():
231 """Test that fallback is used when AI returns invalid JSON."""
232 frames = _make_frames(10)
233 categories = {"code": {"description": "Code editors"}}
234
235 with patch("think.models.generate") as mock_generate:
236 mock_generate.return_value = "not valid json"
237 result = select_frames_for_extraction(
238 frames, max_extractions=5, categories=categories
239 )
240
241 # Should still get a valid result from fallback
242 assert len(result) >= 5
243 assert 1 in result
244
245
246def test_no_ai_selection_without_categories():
247 """Test that AI selection is skipped when categories is None."""
248 frames = _make_frames(10)
249
250 with patch("think.models.generate") as mock_generate:
251 result = select_frames_for_extraction(
252 frames, max_extractions=5, categories=None
253 )
254
255 # generate should not be called
256 mock_generate.assert_not_called()
257 # Should get fallback result
258 assert len(result) >= 5
259 assert 1 in result