From 82f803403cfcea01ceb753b57efa1389c4fd9646 Mon Sep 17 00:00:00 2001 From: Huang Qi Date: Sun, 10 May 2026 12:35:43 +0800 Subject: [PATCH] feat(adapter): implement click-anchored segmentation --- ai_mouse/data_adapters/balabit.py | 36 +++++++++++- tests/test_balabit_adapter.py | 92 +++++++++++++++++++++++++++++++ 2 files changed, 126 insertions(+), 2 deletions(-) diff --git a/ai_mouse/data_adapters/balabit.py b/ai_mouse/data_adapters/balabit.py index c1e41f0..5f422f3 100644 --- a/ai_mouse/data_adapters/balabit.py +++ b/ai_mouse/data_adapters/balabit.py @@ -82,8 +82,40 @@ def segment_by_clicks( window_ms: int, session_id: str, ) -> list[Segment]: - """Stub — implemented in Task 4.""" - raise NotImplementedError + """Extract click-anchored segments from a session. + + For each Left-button Pressed event, collect all Move events within + [click_t - window_ms, click_t) into one segment. + + Args: + events: Full session events (any state, any order is OK but typically sorted). + window_ms: How far back to look before each click. + session_id: String tag attached to every segment for debugging. + + Returns: + List of Segment, one per Left Pressed event that has at least one preceding Move. + """ + segments: list[Segment] = [] + for ev in events: + if ev.state != "Pressed" or ev.button != "Left": + continue + click_t = ev.t_ms + window_start = click_t - window_ms + moves = [ + m for m in events + if m.state == "Move" and window_start <= m.t_ms < click_t + ] + if not moves: + continue + moves.sort(key=lambda m: m.t_ms) + segments.append(Segment( + events=moves, + click_x=ev.x, + click_y=ev.y, + click_t_ms=click_t, + session_id=session_id, + )) + return segments def filter_segments( diff --git a/tests/test_balabit_adapter.py b/tests/test_balabit_adapter.py index 2716740..1bbe8ac 100644 --- a/tests/test_balabit_adapter.py +++ b/tests/test_balabit_adapter.py @@ -93,3 +93,95 @@ class TestParseSessionCsv: csv.write_text("record timestamp,client timestamp,button,state,x,y\n", encoding="utf-8") events = parse_session_csv(csv) assert events == [] + + +class TestSegmentByClicks: + def _make_event(self, t_ms: int, state: str, x: int, y: int, button: str = "NoButton"): + from ai_mouse.data_adapters.balabit import MouseEvent + return MouseEvent(t_ms=t_ms, button=button, state=state, x=x, y=y) + + def test_one_click_one_segment(self): + from ai_mouse.data_adapters.balabit import segment_by_clicks + events = [ + self._make_event(0, "Move", 10, 20), + self._make_event(100, "Move", 50, 60), + self._make_event(500, "Move", 100, 100), + self._make_event(600, "Pressed", 110, 110, button="Left"), + ] + segments = segment_by_clicks(events, window_ms=1200, session_id="test_s1") + assert len(segments) == 1 + seg = segments[0] + assert seg.click_x == 110 + assert seg.click_y == 110 + assert seg.click_t_ms == 600 + assert len(seg.events) == 3 + assert seg.session_id == "test_s1" + + def test_window_excludes_old_events(self): + """Move events earlier than (click_t - window_ms) are dropped.""" + from ai_mouse.data_adapters.balabit import segment_by_clicks + events = [ + self._make_event(0, "Move", 10, 20), # too old + self._make_event(100, "Move", 20, 30), # too old + self._make_event(900, "Move", 30, 40), # in window + self._make_event(1000, "Pressed", 40, 50, button="Left"), + ] + segments = segment_by_clicks(events, window_ms=200, session_id="s") + assert len(segments) == 1 + assert len(segments[0].events) == 1 + assert segments[0].events[0].t_ms == 900 + + def test_multiple_clicks_multiple_segments(self): + from ai_mouse.data_adapters.balabit import segment_by_clicks + events = [ + self._make_event(0, "Move", 10, 20), + self._make_event(100, "Pressed", 50, 50, button="Left"), + self._make_event(200, "Released", 50, 50, button="Left"), + self._make_event(300, "Move", 60, 60), + self._make_event(400, "Move", 70, 70), + self._make_event(500, "Pressed", 80, 80, button="Left"), + ] + segments = segment_by_clicks(events, window_ms=1200, session_id="s") + assert len(segments) == 2 + assert segments[0].click_x == 50 + assert segments[1].click_x == 80 + # Second segment's events must not include the first Pressed + for e in segments[1].events: + assert e.state == "Move" + + def test_skips_pressed_with_non_left_button(self): + """Right-clicks and wheel-clicks don't anchor segments (only Left).""" + from ai_mouse.data_adapters.balabit import segment_by_clicks + events = [ + self._make_event(0, "Move", 10, 20), + self._make_event(100, "Pressed", 50, 50, button="Right"), # ignored + self._make_event(200, "Move", 60, 60), + self._make_event(300, "Pressed", 70, 70, button="Left"), # anchor + ] + segments = segment_by_clicks(events, window_ms=1200, session_id="s") + assert len(segments) == 1 + assert segments[0].click_x == 70 + + def test_no_clicks_returns_empty(self): + from ai_mouse.data_adapters.balabit import segment_by_clicks + events = [ + self._make_event(0, "Move", 10, 20), + self._make_event(100, "Move", 20, 30), + ] + segments = segment_by_clicks(events, window_ms=1200, session_id="s") + assert segments == [] + + def test_excludes_drag_events(self): + """Drag events are not Move; segment should only include Move.""" + from ai_mouse.data_adapters.balabit import segment_by_clicks + events = [ + self._make_event(0, "Move", 10, 20), + self._make_event(100, "Drag", 30, 40), # not Move + self._make_event(200, "Move", 50, 60), + self._make_event(300, "Pressed", 70, 80, button="Left"), + ] + segments = segment_by_clicks(events, window_ms=1200, session_id="s") + assert len(segments) == 1 + # Drag event should not appear in seg.events + assert all(e.state == "Move" for e in segments[0].events) + assert len(segments[0].events) == 2