import sys import unittest from collections import Counter from pathlib import Path MODULE_DIR = Path(__file__).resolve().parents[1] sys.path.insert(0, str(MODULE_DIR)) import export_swe_chat_prose as prose_export def row(session, turn, ordinal, content, **overrides): value = { "session_id": session, "turn_id": turn, "conversation_turn_number": ordinal, "turn_number": ordinal, "turn_type": "user_prompt" if ordinal % 2 == 0 else "assistant_response", "role": "user" if ordinal % 2 == 0 else "assistant", "is_conversational": True, "is_continuation": False, "content": content, } value.update(overrides) return value def select(rows, *, max_per_session=3): """Run both passes the way `export` does, over one in-memory row list.""" markers, funnel = prose_export.conversational_markers(rows) pairs = prose_export.select_turn_ends(markers, funnel, max_per_session=max_per_session) candidates = prose_export.attach_text( rows, pairs, sessions={}, funnel=funnel, character_limit=prose_export.DEFAULT_CHARACTER_LIMIT, ) return candidates, funnel class TurnEndSelectionTests(unittest.TestCase): def test_takes_only_the_last_reply_of_a_multi_message_assistant_run(self): rows = [ row("s1", "t0", 0, "Add the settings view."), row("s1", "t1", 1, "Working on it."), row("s1", "t2", 2, "Still working.", turn_type="assistant_response", role="assistant"), row("s1", "t3", 3, "The settings view is done."), ] candidates, funnel = select(rows) self.assertEqual(1, len(candidates)) self.assertEqual("The settings view is done.", candidates[0].prose) self.assertEqual("t3", candidates[0].pair.response_turn_id) self.assertEqual(2, funnel["rejectedMidTurnAssistantResponse"]) def test_keeps_the_final_reply_when_the_session_ends(self): rows = [ row("s1", "t0", 0, "Ship it."), row("s1", "t1", 1, "Shipped. Anything else?"), ] candidates, _ = select(rows) self.assertEqual(["Shipped. Anything else?"], [c.prose for c in candidates]) def test_pairs_each_reply_with_the_prompt_that_opened_its_own_turn(self): rows = [ row("s1", "t0", 0, "First request."), row("s1", "t1", 1, "First answer."), row("s1", "t2", 2, "Second request."), row("s1", "t3", 3, "Second answer."), ] candidates, _ = select(rows) self.assertEqual( [("First request.", "First answer."), ("Second request.", "Second answer.")], [(c.teacher_prompt, c.prose) for c in candidates], ) def test_orders_by_conversation_ordinal_not_row_arrival(self): rows = [ row("s1", "t3", 3, "Second answer."), row("s1", "t1", 1, "First answer."), row("s1", "t2", 2, "Second request."), row("s1", "t0", 0, "First request."), ] candidates, _ = select(rows) self.assertEqual(["First answer.", "Second answer."], [c.prose for c in candidates]) def test_drops_a_reply_with_no_preceding_user_prompt(self): rows = [row("s1", "t1", 1, "Orphan reply.")] candidates, funnel = select(rows) self.assertEqual([], candidates) self.assertEqual(1, funnel["rejectedNoPrecedingUserPrompt"]) def test_caps_replies_taken_from_one_session(self): rows = [] for index in range(4): rows.append(row("s1", f"u{index}", index * 2, f"Request {index}.")) rows.append(row("s1", f"a{index}", index * 2 + 1, f"Answer {index}.")) candidates, funnel = select(rows, max_per_session=2) self.assertEqual(["Answer 0.", "Answer 1."], [c.prose for c in candidates]) self.assertEqual(2, funnel["rejectedSessionCap"]) def test_rejects_non_conversational_continuation_and_malformed_rows(self): rows = [ row("s1", "t0", 0, "Request."), row("s1", "t1", 1, "Tool output.", is_conversational=False), row("s1", "t2", 2, "Request.", is_continuation=True), row("s1", "t3", 3, "Reply.", conversation_turn_number=None), row("s1", "t4", 4, "Reply.", turn_type="tool_call", role="assistant"), ] _, funnel = select(rows) self.assertEqual(1, funnel["rejectedNonConversational"]) self.assertEqual(1, funnel["rejectedContinuation"]) self.assertEqual(1, funnel["rejectedMalformedMarker"]) self.assertEqual(1, funnel["rejectedTurnTypeOrRole"]) def test_rejects_a_duplicate_turn_id_within_a_session(self): rows = [row("s1", "t0", 0, "a"), row("s1", "t0", 2, "b")] with self.assertRaises(prose_export.DataError): prose_export.conversational_markers(rows) class ProseExtractionTests(unittest.TestCase): def test_student_text_is_the_runtime_extraction_not_the_raw_reply(self): reply = "The migration is done.\n```swift\nstruct View {}\n```\nTests pass." rows = [row("s1", "t0", 0, "Migrate it."), row("s1", "t1", 1, reply)] candidates, _ = select(rows) self.assertEqual("The migration is done.\nTests pass.", candidates[0].prose) self.assertFalse(candidates[0].tail_biased) def test_flags_prose_the_character_limit_actually_truncated(self): rows = [ row("s1", "t0", 0, "Do it."), row("s1", "t1", 1, "old words " * 6 + "Now the settings view."), ] markers, funnel = prose_export.conversational_markers(rows) pairs = prose_export.select_turn_ends(markers, funnel, max_per_session=3) candidates = prose_export.attach_text( rows, pairs, sessions={}, funnel=funnel, character_limit=25 ) self.assertEqual("Now the settings view.", candidates[0].prose) self.assertTrue(candidates[0].tail_biased) def test_drops_a_reply_that_is_entirely_code(self): rows = [ row("s1", "t0", 0, "Show me the struct."), row("s1", "t1", 1, "```swift\nstruct View {}\n```"), ] candidates, funnel = select(rows) self.assertEqual([], candidates) self.assertEqual(1, funnel["rejectedNoProseAfterExtraction"]) def test_drops_empty_and_nul_bearing_content(self): rows = [ row("s1", "t0", 0, "Request."), row("s1", "t1", 1, " "), row("s1", "t2", 2, "Request."), row("s1", "t3", 3, "reply\x00"), ] candidates, funnel = select(rows) self.assertEqual([], candidates) self.assertEqual(2, funnel["rejectedMalformedOrEmptyContent"]) self.assertEqual(2, funnel["rejectedMissingText"]) def test_records_the_question_signal_and_hashes_the_student_text(self): rows = [ row("s1", "t0", 0, "Finish the backend."), row("s1", "t1", 1, "Backend is done. Should I start the UI?"), ] candidates, _ = select(rows) record = candidates[0].json("abc123", prose_export.DEFAULT_CHARACTER_LIMIT) self.assertTrue(record["endsInQuestion"]) self.assertEqual(["t0", "t1"], record["sourceTurnIDs"]) self.assertEqual("abc123", record["sourceRevision"]) self.assertEqual( prose_export.prompt_hash("Backend is done. Should I start the UI?"), record["proseHash"], ) self.assertEqual("Backend is done. Should I start the UI?", record["prose"]) class CapAndDedupeTests(unittest.TestCase): def make(self, prose, repo, user): pair = prose_export.Pair("s", "t0", "t1", 1) return prose_export.Candidate( session_id="s", repo_id=repo, user_id=user, pair=pair, teacher_prompt="Request.", prose=prose, tail_biased=False, ) def test_dedupes_repeated_sign_off_prose_across_sessions(self): candidates, funnel = prose_export.cap_and_dedupe( [ self.make("All tests pass.", "r1", "u1"), self.make(" all TESTS pass. ", "r2", "u2"), self.make("Renamed the module.", "r3", "u3"), ], Counter(), max_per_repo=10, max_per_user=10, ) self.assertEqual(["All tests pass.", "Renamed the module."], [c.prose for c in candidates]) self.assertEqual(1, funnel["rejectedDuplicateProse"]) def test_caps_repository_and_user_concentration(self): candidates, funnel = prose_export.cap_and_dedupe( [ self.make("One.", "r1", "u1"), self.make("Two.", "r1", "u2"), self.make("Three.", "r2", "u1"), self.make("Four.", "r2", "u2"), ], Counter(), max_per_repo=1, max_per_user=1, ) self.assertEqual(["One.", "Four."], [c.prose for c in candidates]) self.assertEqual(1, funnel["rejectedRepoCap"]) self.assertEqual(1, funnel["rejectedUserCap"]) self.assertEqual(2, funnel["exportedCandidates"]) class ArgumentTests(unittest.TestCase): def test_rejects_a_mutable_revision(self): for revision in ["main", ""]: with self.subTest(revision), self.assertRaises(prose_export.DataError): prose_export.export( conversations=[], sessions_path=[], revision=revision, output=Path("/dev/null"), manifest_path=Path("/dev/null"), max_per_repo=1, max_per_user=1, max_per_session=1, character_limit=1, ) def test_rejects_non_positive_caps_and_limits(self): for kwargs in [ {"max_per_repo": 0}, {"max_per_user": 0}, {"max_per_session": 0}, {"character_limit": 0}, ]: settings = { "max_per_repo": 1, "max_per_user": 1, "max_per_session": 1, "character_limit": 1, **kwargs, } with self.subTest(kwargs), self.assertRaises(prose_export.DataError): prose_export.export( conversations=[], sessions_path=[], revision="abc123", output=Path("/dev/null"), manifest_path=Path("/dev/null"), **settings, ) if __name__ == "__main__": unittest.main()