nucleic-purpose-classifier/tests/test_export_swe_chat.py

import sys
import unittest
from pathlib import Path


MODULE_DIR = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(MODULE_DIR))

import export_swe_chat


def row(session, turn, prompt, **overrides):
    value = {
        "session_id": session,
        "turn_id": turn,
        "conversation_turn_number": int(turn[1:]),
        "turn_number": int(turn[1:]),
        "turn_type": "user_prompt",
        "role": "user",
        "is_conversational": True,
        "is_continuation": False,
        "content": prompt,
    }
    value.update(overrides)
    return value


class ExportSWEChatTests(unittest.TestCase):
    def test_selects_first_prompt_and_its_response_dedupes_and_caps_sources(self):
        rows = [
            row("s1", "t2", "answer", conversation_turn_number=1, turn_type="assistant_response", role="assistant"),
            row("s1", "t1", "first", conversation_turn_number=0),
            row("s2", "t1", " first ", conversation_turn_number=0),
            row("s2", "t2", "answer", conversation_turn_number=1, turn_type="assistant_response", role="assistant"),
            row("s3", "t1", "one"),
            row("s4", "t1", "another one", conversation_turn_number=0),
            row("s4", "t2", "answer", conversation_turn_number=1, turn_type="assistant_response", role="assistant"),
        ]
        candidates, funnel = export_swe_chat.select_candidates(
            rows,
            sessions={"s1": ("repo", "user"), "s2": ("repo", "user"), "s4": ("repo", "user")},
            max_per_repo=1, max_per_user=1,
        )
        self.assertEqual(1, len(candidates))
        self.assertEqual(["t1", "t2"], [candidates[0].prompt_turn.turn_id, candidates[0].response_turn.turn_id])
        self.assertEqual("first", candidates[0].prompt_turn.prompt)
        self.assertEqual(1, funnel["rejectedDuplicateFirstPrompt"])
        self.assertEqual(1, funnel["sessionsWithoutFirstAssistantResponse"])
        self.assertEqual(1, funnel["rejectedRepoCap"])

    def test_filters_non_user_continuation_empty_and_missing_response(self):
        rows = [
            row("s1", "t1", "x", role="assistant"), row("s1", "t2", "x", is_continuation=True), row("s1", "t3", ""),
            row("s2", "t1", "a", conversation_turn_number=0), row("s2", "t2", "wrong ordinal", conversation_turn_number=2, turn_type="assistant_response", role="assistant"),
        ]
        candidates, funnel = export_swe_chat.select_candidates(rows, sessions={}, max_per_repo=10, max_per_user=10)
        self.assertEqual([], candidates)
        self.assertEqual(1, funnel["rejectedRole"])
        self.assertEqual(1, funnel["rejectedContinuation"])
        self.assertEqual(1, funnel["rejectedUserMalformedOrEmpty"])
        self.assertEqual(1, funnel["sessionsWithoutFirstAssistantResponse"])

    def test_candidate_only_retains_first_prompt_as_student_text(self):
        prompt = export_swe_chat.Turn("s", "t1", 0, 0, "first")
        response = export_swe_chat.Turn("s", "t2", 1, 1, "answer")
        value = export_swe_chat.Candidate("s", "r", "u", prompt, response).json("a" * 40)
        self.assertEqual("first", value["prompt"])
        self.assertEqual("answer", value["teacherResponse"])
        self.assertNotIn("answer", value["teacherResponseHash"])


if __name__ == "__main__":
    unittest.main()