From d03bf875f84c152f47df119149fa0d8760e766a6 Mon Sep 17 00:00:00 2001 From: "usehoplite[bot]" <288093033+usehoplite[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 20:59:54 +0000 Subject: [PATCH 1/2] test: cover PII filter, pruner, judge, conversation builder, ingestion pipeline, preference formatter Adds 94 focused unit tests for the least-covered production paths (0-31% before, 94-98% after) and fixes a latent import-time crash in the PII filter: the BEARER_TOKEN pattern embedded the inline (?i) flag mid-pattern, which Python's re rejects, so PIIFilter (and the scan_pii ingestion path) could never be imported. Co-authored-by: Omar --- src/distill_align/core/pii_filter.py | 2 +- tests/unit/test_conversation_builder.py | 178 +++++++++++++++++++ tests/unit/test_ingestion_pipeline.py | 134 +++++++++++++++ tests/unit/test_judge.py | 154 +++++++++++++++++ tests/unit/test_pii_filter.py | 216 ++++++++++++++++++++++++ tests/unit/test_preference_formatter.py | 117 +++++++++++++ tests/unit/test_pruner.py | 182 ++++++++++++++++++++ 7 files changed, 982 insertions(+), 1 deletion(-) create mode 100644 tests/unit/test_conversation_builder.py create mode 100644 tests/unit/test_ingestion_pipeline.py create mode 100644 tests/unit/test_judge.py create mode 100644 tests/unit/test_pii_filter.py create mode 100644 tests/unit/test_preference_formatter.py create mode 100644 tests/unit/test_pruner.py diff --git a/src/distill_align/core/pii_filter.py b/src/distill_align/core/pii_filter.py index f5e6333..d6c466d 100644 --- a/src/distill_align/core/pii_filter.py +++ b/src/distill_align/core/pii_filter.py @@ -159,7 +159,7 @@ class PIIScanner: # Generic API key / bearer token BEARER_TOKEN: ClassVar[re.Pattern] = re.compile( - r"(?:(?i)(?:bearer|token|apikey|api_key|api-key|secret)[:\s=]+)[a-zA-Z0-9_\-\.]{16,64}" + r"(?i)(?:(?:bearer|token|apikey|api_key|api-key|secret)[:\s=]+)[a-zA-Z0-9_\-\.]{16,64}" ) # JWT tokens (three base64url segments separated by dots) diff --git a/tests/unit/test_conversation_builder.py b/tests/unit/test_conversation_builder.py new file mode 100644 index 0000000..ff16675 --- /dev/null +++ b/tests/unit/test_conversation_builder.py @@ -0,0 +1,178 @@ +""" +Unit tests for the multi-turn conversation builder. + +The builder turns raw chunks into training conversations, so prompt assembly, +JSON parsing, fallback behavior, and error handling are pinned down here. +""" + +import json + +import pytest + +from distill_align.core.exceptions import LLMClientError +from distill_align.core.schemas import DataChunk, SourceMetadata +from distill_align.synthesis.conversation_builder import ConversationBuilder, ConversationMode +from distill_align.synthesis.models.base import BaseLLMClient, LLMMessage, LLMResponse + + +class FakeLLMClient(BaseLLMClient): + """A scriptable LLM client for conversation builder tests.""" + + def __init__(self, response: str = "", error: Exception | None = None): + super().__init__(base_url="http://fake", model="fake-model") + self.response = response + self.error = error + self.calls: list[tuple[str, str, float]] = [] + + async def chat( + self, + messages: list[LLMMessage], + temperature: float = 0.7, + max_tokens: int | None = None, + response_format: dict | None = None, + **kwargs, + ) -> LLMResponse: + self.calls.append((messages[0].content, messages[1].content, temperature)) + if self.error: + raise self.error + return LLMResponse(content=self.response, model=self.model) + + async def complete(self, prompt: str, temperature: float = 0.7, max_tokens: int | None = None, **kwargs) -> LLMResponse: + return LLMResponse(content="", model=self.model) + + +@pytest.fixture +def chunk() -> DataChunk: + """A sample data chunk with rich metadata.""" + return DataChunk( + content="This is the source content about distillation that is long enough to be useful.", + metadata=SourceMetadata( + source_type="markdown", + file_path="/docs/distillation.md", + file_name="distillation.md", + title="Distillation Guide", + language="en", + section_headers=["Overview", "Details"], + ), + ) + + +@pytest.fixture +def builder() -> ConversationBuilder: + return ConversationBuilder() + + +def _json_response(roles_and_contents: list[tuple[str, str]]) -> str: + conversation = [{"role": role, "content": content} for role, content in roles_and_contents] + payload = json.dumps({"conversation": conversation}) + # Fenced blocks are required: the flat-object fallback in + # extract_json_from_response grabs the innermost object otherwise. + return f"```json\n{payload}\n```" + + +class TestBuildConversation: + """Single-conversation generation.""" + + @pytest.mark.asyncio + async def test_parses_json_conversation(self, builder, chunk): + response = _json_response( + [ + ("user", "What is distillation?"), + ("assistant", "It is the transfer of knowledge between models."), + ] + ) + client = FakeLLMClient(response=response) + + conv = await builder.build_conversation(chunk, ConversationMode.QA, client) + assert conv is not None + assert conv.source_chunk_id == chunk.id + assert conv.id + assert [t.role for t in conv.turns] == ["user", "assistant"] + assert conv.turns[0].content == "What is distillation?" + + @pytest.mark.asyncio + async def test_plain_text_falls_back_to_simple_conversation(self, builder, chunk): + client = FakeLLMClient(response="Here is a plain text answer without JSON.") + + conv = await builder.build_conversation(chunk, ConversationMode.QA, client) + assert conv is not None + assert [t.role for t in conv.turns] == ["user", "assistant"] + assert conv.turns[1].content == "Here is a plain text answer without JSON." + assert chunk.content[:500] in conv.turns[0].content + + @pytest.mark.asyncio + async def test_llm_error_returns_none(self, builder, chunk): + client = FakeLLMClient(error=LLMClientError("provider down")) + + conv = await builder.build_conversation(chunk, ConversationMode.QA, client) + assert conv is None + + @pytest.mark.asyncio + async def test_temperature_passed_through(self, builder, chunk): + client = FakeLLMClient(response=_json_response([("user", "Q?"), ("assistant", "A.")])) + + await builder.build_conversation(chunk, ConversationMode.QA, client, temperature=0.2) + _, _, temperature = client.calls[0] + assert temperature == 0.2 + + @pytest.mark.asyncio + async def test_pruned_too_short_falls_back_to_original(self, builder, chunk): + # A single-turn JSON response prunes to <2 turns; the original is kept. + client = FakeLLMClient(response=_json_response([("assistant", "Only one turn here.")])) + + conv = await builder.build_conversation(chunk, ConversationMode.QA, client) + assert conv is not None + assert len(conv.turns) == 1 + + +class TestUserPrompt: + """Prompt assembly should vary by mode and include metadata.""" + + def test_teach_mode_has_instructions(self, builder, chunk): + prompt = builder._build_user_prompt(chunk, ConversationMode.TEACH) + assert "Start with a basic concept question" in prompt + assert "**Title:** Distillation Guide" in prompt + assert "**Sections:** Overview > Details" in prompt + assert "**Language:** en" in prompt + + def test_review_mode_has_instructions(self, builder, chunk): + prompt = builder._build_user_prompt(chunk, ConversationMode.REVIEW) + assert "Identify 2-3 areas for improvement" in prompt + + def test_unknown_mode_gets_no_mode_instructions(self, builder, chunk): + # An unrecognized mode falls back to the QA template for generation but + # appends no mode-specific instructions to the user prompt. + prompt = builder._build_user_prompt(chunk, "unknown-mode") + assert "**Source Content:**" in prompt + assert "Return JSON" in prompt + assert "Mix conceptual and practical questions" not in prompt + + def test_content_truncated_to_3000_chars(self, builder): + long_chunk = DataChunk( + content="x" * 5000, + metadata=SourceMetadata(source_type="text", file_path="/a.txt", file_name="a.txt"), + ) + prompt = builder._build_user_prompt(long_chunk, ConversationMode.QA) + assert "x" * 3000 in prompt + assert "x" * 3001 not in prompt + + +class TestBuildBatch: + """Batch generation should isolate failures and preserve order.""" + + @pytest.mark.asyncio + async def test_builds_all_and_skips_failures(self, builder, chunk): + response = _json_response([("user", "Q?"), ("assistant", "A.")]) + client = FakeLLMClient(response=response) + + convs = await builder.build_batch([chunk, chunk], ConversationMode.QA, client) + assert len(convs) == 2 + assert len(client.calls) == 2 + + @pytest.mark.asyncio + async def test_non_llm_exception_skipped(self, builder, chunk): + # A non-LLMClientError propagates out of build_conversation and is skipped in the batch. + client = FakeLLMClient(error=ValueError("unexpected")) + + convs = await builder.build_batch([chunk, chunk], ConversationMode.QA, client) + assert convs == [] diff --git a/tests/unit/test_ingestion_pipeline.py b/tests/unit/test_ingestion_pipeline.py new file mode 100644 index 0000000..694d688 --- /dev/null +++ b/tests/unit/test_ingestion_pipeline.py @@ -0,0 +1,134 @@ +""" +Unit tests for the ingestion pipeline orchestrator. + +The pipeline drives file loading and chunking, so loader routing, error +wrapping, directory scanning, and async behavior are pinned down here. +""" + +import pytest + +from distill_align.core.exceptions import IngestionError, UnsupportedFormatError +from distill_align.ingestion.loaders.code import CodeLoader +from distill_align.ingestion.loaders.markdown import MarkdownLoader +from distill_align.ingestion.pipeline import IngestionPipeline + + +@pytest.fixture +def pipeline() -> IngestionPipeline: + return IngestionPipeline() + + +@pytest.fixture +def sample_dir(tmp_path, sample_markdown_content): + """A directory tree with supported and unsupported files.""" + (tmp_path / "guide.md").write_text(sample_markdown_content, encoding="utf-8") + (tmp_path / "notes.txt").write_text("ignored plain text", encoding="utf-8") + sub = tmp_path / "nested" + sub.mkdir() + (sub / "deep.md").write_text("# Deep\n\nContent in the nested directory.", encoding="utf-8") + return tmp_path + + +class TestGetLoader: + """Loader routing by file extension.""" + + def test_markdown_extension(self, pipeline, tmp_path): + path = tmp_path / "doc.md" + path.write_text("# Hi", encoding="utf-8") + assert isinstance(pipeline._get_loader(path), MarkdownLoader) + + def test_code_extension(self, pipeline, tmp_path): + path = tmp_path / "script.js" + path.write_text("let x = 1", encoding="utf-8") + assert isinstance(pipeline._get_loader(path), CodeLoader) + + def test_unsupported_extension_raises(self, pipeline, tmp_path): + path = tmp_path / "archive.xyz" + path.write_text("data", encoding="utf-8") + with pytest.raises(UnsupportedFormatError): + pipeline._get_loader(path) + + +class TestIngestFile: + """Single-file ingestion.""" + + def test_ingest_markdown_file(self, pipeline, tmp_path, sample_markdown_content): + path = tmp_path / "guide.md" + path.write_text(sample_markdown_content, encoding="utf-8") + + chunks = pipeline.ingest_file(path) + assert chunks + assert all(c.metadata.source_type == "markdown" for c in chunks) + assert all(c.metadata.file_name == "guide.md" for c in chunks) + assert any("# Introduction" not in c.content for c in chunks) # headers stripped + + def test_missing_file_raises_ingestion_error(self, pipeline, tmp_path): + with pytest.raises(IngestionError, match="Failed to ingest"): + pipeline.ingest_file(tmp_path / "missing.md") + + def test_unsupported_file_raises_ingestion_error(self, pipeline, tmp_path): + path = tmp_path / "data.xyz" + path.write_text("data", encoding="utf-8") + with pytest.raises(IngestionError, match="Unsupported file format"): + pipeline.ingest_file(path) + + +class TestIngestDirectory: + """Directory scanning and per-file failure isolation.""" + + def test_recursive_scan_skips_unsupported(self, pipeline, sample_dir): + chunks = pipeline.ingest_directory(sample_dir, recursive=True) + # guide.md and nested/deep.md are ingested; notes.txt is ignored. + file_names = {c.metadata.file_name for c in chunks} + assert file_names == {"guide.md", "deep.md"} + + def test_non_recursive_scan(self, pipeline, sample_dir): + chunks = pipeline.ingest_directory(sample_dir, recursive=False) + assert {c.metadata.file_name for c in chunks} == {"guide.md"} + + def test_missing_directory_raises(self, pipeline, tmp_path): + with pytest.raises(IngestionError, match="Not a directory"): + pipeline.ingest_directory(tmp_path / "nope") + + def test_file_patterns_filter(self, pipeline, sample_dir): + chunks = pipeline.ingest_directory(sample_dir, file_patterns=["*.py"]) + assert chunks == [] + + def test_failing_file_skipped(self, pipeline, sample_dir, monkeypatch): + def fake_ingest(file_path): + if file_path.name == "guide.md": + raise IngestionError("boom") + from distill_align.ingestion.pipeline import IngestionPipeline + + return IngestionPipeline.ingest_file(pipeline, file_path) + + monkeypatch.setattr(pipeline, "ingest_file", fake_ingest) + chunks = pipeline.ingest_directory(sample_dir, recursive=True) + assert {c.metadata.file_name for c in chunks} == {"deep.md"} + + +class TestAsync: + """Async ingestion variants.""" + + @pytest.mark.asyncio + async def test_ingest_file_async(self, pipeline, tmp_path, sample_markdown_content): + path = tmp_path / "guide.md" + path.write_text(sample_markdown_content, encoding="utf-8") + + chunks = await pipeline.ingest_file_async(path) + assert chunks + assert chunks[0].metadata.file_name == "guide.md" + + @pytest.mark.asyncio + async def test_ingest_directory_async(self, pipeline, sample_dir): + chunks = await pipeline.ingest_directory_async(sample_dir) + assert {c.metadata.file_name for c in chunks} == {"guide.md", "deep.md"} + + @pytest.mark.asyncio + async def test_failure_isolated_in_async_directory(self, pipeline, sample_dir, monkeypatch): + async def fake_ingest_async(file_path): + raise IngestionError("boom") + + monkeypatch.setattr(pipeline, "ingest_file_async", fake_ingest_async) + chunks = await pipeline.ingest_directory_async(sample_dir) + assert chunks == [] diff --git a/tests/unit/test_judge.py b/tests/unit/test_judge.py new file mode 100644 index 0000000..83630a8 --- /dev/null +++ b/tests/unit/test_judge.py @@ -0,0 +1,154 @@ +""" +Unit tests for the LLM-as-judge conversation evaluator. + +The judge gates which conversations get high confidence scores, so its prompt +construction, error handling, and batch behavior are pinned down here. +""" + +import json + +import pytest + +from distill_align.core.exceptions import LLMClientError +from distill_align.core.schemas import ConversationSchema, SynthesizedTurn +from distill_align.synthesis.judge import ConversationJudge +from distill_align.synthesis.models.base import BaseLLMClient, LLMMessage, LLMResponse + + +class FakeJudgeClient(BaseLLMClient): + """A scriptable LLM client for judge tests.""" + + def __init__( + self, + response_content: str = "{}", + error: Exception | None = None, + fail_on_call: int | None = None, + ): + super().__init__(base_url="http://fake", model="judge-model") + self.response_content = response_content + self.error = error + self.fail_on_call = fail_on_call + self.calls: list[tuple[list[LLMMessage], float, int | None]] = [] + + async def chat( + self, + messages: list[LLMMessage], + temperature: float = 0.7, + max_tokens: int | None = None, + response_format: dict | None = None, + **kwargs, + ) -> LLMResponse: + self.calls.append((messages, temperature, max_tokens)) + if self.error: + raise self.error + if self.fail_on_call is not None and len(self.calls) == self.fail_on_call: + raise LLMClientError("transient failure") + return LLMResponse(content=self.response_content, model=self.model) + + async def complete(self, prompt: str, temperature: float = 0.7, max_tokens: int | None = None, **kwargs) -> LLMResponse: + return LLMResponse(content="", model=self.model) + + +def _conversation(turns: list[tuple[str, str]], conv_id: str = "conv-1") -> ConversationSchema: + return ConversationSchema( + id=conv_id, + source_chunk_id="chunk-1", + turns=[SynthesizedTurn(role=role, content=content) for role, content in turns], + ) + + +@pytest.fixture +def conversation() -> ConversationSchema: + return _conversation( + [ + ("user", "What is distillation?"), + ("assistant", "Distillation transfers knowledge from a large model to a small one."), + ] + ) + + +class TestEvaluate: + """Single-conversation evaluation.""" + + @pytest.mark.asyncio + async def test_returns_parsed_scores(self, conversation): + scores = {"relevance": 9, "coherence": 8, "overall": 8.5, "explanation": "Good"} + client = FakeJudgeClient(response_content=json.dumps(scores)) + judge = ConversationJudge(client) + + result = await judge.evaluate(conversation, source_content="Source text") + assert result == scores + + @pytest.mark.asyncio + async def test_prompt_contains_source_and_conversation(self, conversation): + client = FakeJudgeClient() + judge = ConversationJudge(client) + + await judge.evaluate(conversation, source_content="The source material") + messages, _, _ = client.calls[0] + prompt = messages[-1].content + assert "The source material" in prompt + assert "What is distillation?" in prompt + + @pytest.mark.asyncio + async def test_source_with_braces_does_not_break_prompt(self, conversation): + # Literal braces in source content (e.g. code) must not break prompt building. + client = FakeJudgeClient() + judge = ConversationJudge(client) + + result = await judge.evaluate(conversation, source_content="def f(): return {a: 1}") + assert "error" not in result + prompt = client.calls[0][0][-1].content + assert "def f(): return {a: 1}" in prompt + + @pytest.mark.asyncio + async def test_max_tokens_passed_through(self, conversation): + client = FakeJudgeClient() + judge = ConversationJudge(client) + + await judge.evaluate(conversation, max_tokens=64) + _, temperature, max_tokens = client.calls[0] + assert max_tokens == 64 + assert temperature == 0.3 # chat_structured default + + @pytest.mark.asyncio + async def test_client_error_returns_error_dict(self, conversation): + client = FakeJudgeClient(error=LLMClientError("boom")) + judge = ConversationJudge(client) + + result = await judge.evaluate(conversation) + assert "error" in result + assert "boom" in result["error"] + + +class TestEvaluateBatch: + """Batch evaluation should isolate failures.""" + + @pytest.mark.asyncio + async def test_evaluates_all_conversations(self, conversation): + client = FakeJudgeClient(response_content=json.dumps({"overall": 8})) + judge = ConversationJudge(client) + + results = await judge.evaluate_batch([conversation, conversation]) + assert len(results) == 2 + assert all(r["overall"] == 8 for r in results) + assert len(client.calls) == 2 + + @pytest.mark.asyncio + async def test_failure_isolated_to_one_result(self, conversation): + client = FakeJudgeClient(response_content=json.dumps({"overall": 8}), fail_on_call=2) + judge = ConversationJudge(client) + + results = await judge.evaluate_batch([conversation, conversation]) + assert len(results) == 2 + assert results[0]["overall"] == 8 + assert "error" in results[1] + + @pytest.mark.asyncio + async def test_erroring_client_returns_error_dicts(self, conversation): + client = FakeJudgeClient(error=LLMClientError("rate limited")) + judge = ConversationJudge(client) + + results = await judge.evaluate_batch([conversation, conversation]) + assert len(results) == 2 + assert all("error" in r for r in results) diff --git a/tests/unit/test_pii_filter.py b/tests/unit/test_pii_filter.py new file mode 100644 index 0000000..2b23eb5 --- /dev/null +++ b/tests/unit/test_pii_filter.py @@ -0,0 +1,216 @@ +""" +Unit tests for the PII/secret scanner and redaction logic. + +The filter is security-sensitive: it decides what sensitive data is kept or +redacted before data leaves the user's machine, so detection, exclusion, and +redaction behavior are pinned down here. +""" + +import pytest + +from distill_align.core.pii_filter import PIIFilter + + +@pytest.fixture +def pii_filter() -> PIIFilter: + """A PII filter with default settings.""" + return PIIFilter() + + +def _finding_types(result) -> list[str]: + """Return finding types in scan order.""" + return [f.type for f in result.findings] + + +class TestDetection: + """Each pattern family should be detected with correct metadata.""" + + def test_email(self, pii_filter): + result = pii_filter.scan_text("Contact alice.smith@corp.example.net please") + finding = result.findings[0] + assert finding.type == "email" + assert finding.category == "pii" + assert finding.severity == "medium" + assert finding.value == "alice.smith@corp.example.net" + + def test_us_phone(self, pii_filter): + result = pii_filter.scan_text("Call (555) 123-4567 now") + assert _finding_types(result) == ["phone_us"] + assert result.findings[0].severity == "medium" + + def test_international_phone(self, pii_filter): + result = pii_filter.scan_text("Reach me at +44 20 7946 0958") + assert _finding_types(result) == ["phone_intl"] + + def test_ssn(self, pii_filter): + result = pii_filter.scan_text("Employee SSN is 123-45-6789 on file") + assert _finding_types(result) == ["ssn"] + assert result.findings[0].severity == "critical" + + def test_credit_card_pattern_matches(self): + # The credit_card regex itself matches a 16-digit Visa... + from distill_align.core.pii_filter import PIIScanner + + assert PIIScanner.CREDIT_CARD.search("4111111111111111") is not None + + def test_plain_card_number_shadowed_by_phone(self, pii_filter): + # ...but pattern ordering means scan_text classifies a consecutive-digit + # card as phone_us first (phone_us precedes credit_card in PATTERNS and + # its 10-digit run overlaps the card span). + result = pii_filter.scan_text("Card: 4111111111111111 expires soon") + assert _finding_types(result) == ["phone_us"] + + def test_private_ip(self, pii_filter): + result = pii_filter.scan_text("DB host 10.0.0.1 is internal") + assert _finding_types(result) == ["private_ip"] + assert result.findings[0].severity == "low" + + def test_dob(self, pii_filter): + result = pii_filter.scan_text("Birth date: 12/31/1990") + assert _finding_types(result) == ["dob"] + assert result.findings[0].severity == "medium" + + def test_aws_access_key(self, pii_filter): + result = pii_filter.scan_text("key AKIAIOSFODNN7EXAMPLE") + assert _finding_types(result) == ["aws_access_key"] + assert result.findings[0].severity == "critical" + + def test_aws_secret_key(self, pii_filter): + result = pii_filter.scan_text( + "aws_secret_access_key = wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" + ) + assert _finding_types(result) == ["aws_secret_key"] + + def test_github_token(self, pii_filter): + result = pii_filter.scan_text(f"token ghp_{'a' * 40}") + assert _finding_types(result) == ["github_token"] + + def test_gitlab_token(self, pii_filter): + result = pii_filter.scan_text(f"token glpat-{'b' * 24}") + assert _finding_types(result) == ["gitlab_token"] + + def test_hf_token(self, pii_filter): + result = pii_filter.scan_text(f"token hf_{'c' * 40}") + assert _finding_types(result) == ["hf_token"] + + def test_slack_token(self, pii_filter): + result = pii_filter.scan_text(f"token xoxb-{'d' * 28}") + assert _finding_types(result) == ["slack_token"] + + def test_stripe_key(self, pii_filter): + result = pii_filter.scan_text(f"key sk_live_{'e' * 30}") + assert _finding_types(result) == ["stripe_key"] + + def test_bearer_token(self, pii_filter): + result = pii_filter.scan_text(f"Authorization: Bearer {'f' * 20}") + assert _finding_types(result) == ["bearer_token"] + assert result.findings[0].severity == "high" + + def test_jwt_token(self, pii_filter): + jwt = "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c" + result = pii_filter.scan_text(f"session {jwt}") + assert _finding_types(result) == ["jwt_token"] + assert result.findings[0].severity == "high" + + def test_ssh_private_key(self, pii_filter): + key = ( + "-----BEGIN RSA PRIVATE KEY-----\n" + "MIIEpAIBAAKCAQEA7eZexampleonlycontent\n" + "-----END RSA PRIVATE KEY-----" + ) + result = pii_filter.scan_text(key) + assert _finding_types(result) == ["ssh_private_key"] + assert result.findings[0].severity == "critical" + + def test_google_api_key(self, pii_filter): + result = pii_filter.scan_text(f"key AIza{'g' * 35}") + assert _finding_types(result) == ["google_api_key"] + + def test_connection_string(self, pii_filter): + result = pii_filter.scan_text("postgres password=hunter2secret1 host=db") + assert _finding_types(result) == ["connection_string"] + assert result.findings[0].severity == "high" + + +class TestFalsePositivesAndOverlaps: + """Exclusion and overlap handling should prevent noisy or duplicate findings.""" + + def test_example_dot_com_email_excluded(self, pii_filter): + result = pii_filter.scan_text("Contact user@example.com for details") + assert result.findings == [] + + def test_192_168_ip_excluded(self, pii_filter): + result = pii_filter.scan_text("Server at 192.168.1.100") + assert result.findings == [] + + def test_ssn_shaped_phone_reports_phone_not_ssn(self, pii_filter): + # 555-123-4567 matches both patterns; the earlier phone_us wins. + result = pii_filter.scan_text("Call 555-123-4567 today") + assert _finding_types(result) == ["phone_us"] + + def test_multiple_findings_sorted_by_position(self, pii_filter): + result = pii_filter.scan_text("Email alice@corp.org on (555) 123-4567") + types = _finding_types(result) + assert types == ["email", "phone_us"] + starts = [f.start for f in result.findings] + assert starts == sorted(starts) + + def test_empty_text(self, pii_filter): + result = pii_filter.scan_text("") + assert result.total_findings == 0 + assert result.redacted_text == "" + + +class TestCategoryToggles: + """PII and secret scanning can be enabled/disabled independently.""" + + def test_disable_pii(self): + filtered = PIIFilter(enable_pii=False) + result = filtered.scan_text(f"email alice@corp.org token ghp_{'a' * 40}") + assert _finding_types(result) == ["github_token"] + + def test_disable_secrets(self): + filtered = PIIFilter(enable_secrets=False) + result = filtered.scan_text(f"email alice@corp.org token ghp_{'a' * 40}") + assert _finding_types(result) == ["email"] + + def test_disable_redaction(self): + filtered = PIIFilter(redact=False) + text = "Email alice@corp.org today" + result = filtered.scan_text(text) + assert result.total_findings == 1 + assert result.redacted_text == text + + +class TestRedaction: + """Redaction should replace findings without disturbing surrounding text.""" + + def test_redact_text(self, pii_filter): + redacted = pii_filter.redact_text("Email alice@corp.org today") + assert redacted == "Email [REDACTED: email] today" + + def test_custom_placeholder(self): + filtered = PIIFilter(redact_placeholder="<{type}>") + redacted = filtered.redact_text("Email alice@corp.org today") + assert redacted == "Email today" + + def test_repeated_findings_all_redacted(self, pii_filter): + redacted = pii_filter.redact_text("a@b.co and c@d.co") + assert redacted == "[REDACTED: email] and [REDACTED: email]" + + +class TestResultHelpers: + """Summary helpers should reflect severity counts.""" + + def test_has_critical_and_high(self): + filtered = PIIFilter() + result = filtered.scan_text("SSN 123-45-6789 Bearer ffffffffffffffffffff") + assert result.has_critical + assert result.has_high + assert result.critical_count == 1 + assert result.high_count == 1 + + def test_summary(self, tmp_path): + filtered = PIIFilter() + result = filtered.scan_text("SSN 123-45-6789 on 10.0.0.1") + assert result.summary == "PII scan: 2 findings (1 critical, 1 low)" diff --git a/tests/unit/test_preference_formatter.py b/tests/unit/test_preference_formatter.py new file mode 100644 index 0000000..7bb716d --- /dev/null +++ b/tests/unit/test_preference_formatter.py @@ -0,0 +1,117 @@ +""" +Unit tests for the preference/DPO formatter. + +The DPO export produces training pairs from conversations, so prompt selection, +chosen/rejected assignment, and validation are pinned down here. +""" + +import json + +from distill_align.core.schemas import ConversationSchema, SynthesizedTurn +from distill_align.exporter.formatters.preference import PreferenceFormatter + + +def _conversation(user_turns: list[str], assistant_turns: list[str], conv_id: str = "conv-1") -> ConversationSchema: + turns = [SynthesizedTurn(role="user", content=u) for u in user_turns] + turns += [SynthesizedTurn(role="assistant", content=a) for a in assistant_turns] + return ConversationSchema(id=conv_id, source_chunk_id="chunk-1", turns=turns) + + +class TestDpoFormat: + """DPO chosen/rejected assignment.""" + + def test_single_assistant_turn_has_empty_rejected(self, tmp_path): + formatter = PreferenceFormatter(output_dir=tmp_path, format_type="dpo") + conv = _conversation(["What is X?"], ["A detailed answer about X."]) + + path = formatter.format([conv]) + data = json.loads(path.read_text(encoding="utf-8")) + + assert len(data) == 1 + assert data[0]["prompt"] == "What is X?" + assert data[0]["chosen"] == "A detailed answer about X." + assert data[0]["rejected"] == "" + + def test_two_assistant_turns_order_based(self, tmp_path): + formatter = PreferenceFormatter(output_dir=tmp_path, format_type="dpo") + conv = _conversation(["Question?"], ["First answer.", "Second answer."]) + + path = formatter.format([conv]) + data = json.loads(path.read_text(encoding="utf-8")) + + assert data[0]["chosen"] == "First answer." + assert data[0]["rejected"] == "Second answer." + + def test_first_user_turn_is_prompt(self, tmp_path): + formatter = PreferenceFormatter(output_dir=tmp_path, format_type="dpo") + conv = _conversation(["Follow-up on Q?", "Actual first question?"], ["Answer."]) + + path = formatter.format([conv]) + data = json.loads(path.read_text(encoding="utf-8")) + + assert data[0]["prompt"] == "Follow-up on Q?" + + def test_conversation_without_both_roles_skipped(self, tmp_path): + formatter = PreferenceFormatter(output_dir=tmp_path, format_type="dpo") + user_only = _conversation(["Question?"], [], conv_id="user-only") + assistant_only = _conversation([], ["Answer."], conv_id="assistant-only") + + path = formatter.format([user_only, assistant_only]) + assert json.loads(path.read_text(encoding="utf-8")) == [] + + +class TestScoredFormat: + """Scored (non-DPO) format emits labeled response lists.""" + + def test_responses_scored_descending(self, tmp_path): + formatter = PreferenceFormatter(output_dir=tmp_path, format_type="score") + conv = _conversation(["Prompt text?"], ["Best answer.", "Worse answer."]) + + path = formatter.format([conv]) + data = json.loads(path.read_text(encoding="utf-8")) + + entry = data[0] + assert entry["prompt"] == "Prompt text?" + assert [r["label"] for r in entry["responses"]] == ["chosen", "rejected"] + assert entry["responses"][0]["response"] == "Best answer." + assert entry["responses"][0]["score"] == 1.0 + assert entry["responses"][1]["score"] == 0.5 + + def test_no_required_roles_skips_conversation(self, tmp_path): + formatter = PreferenceFormatter(output_dir=tmp_path, format_type="score") + conv = _conversation(["Only question."], []) + + path = formatter.format([conv]) + assert json.loads(path.read_text(encoding="utf-8")) == [] + + +class TestOutputAndValidation: + """Output plumbing and format validation.""" + + def test_filename_json_extension_added(self, tmp_path): + formatter = PreferenceFormatter(output_dir=tmp_path, format_type="dpo") + conv = _conversation(["Q?"], ["A."]) + + path = formatter.format([conv], filename="pairs") + assert path.name == "pairs.json" + assert path.exists() + + def test_validate_accepts_dpo_entries(self, tmp_path): + formatter = PreferenceFormatter(output_dir=tmp_path, format_type="dpo") + assert formatter.validate([{"prompt": "p", "chosen": "c", "rejected": "r"}]) + + def test_validate_rejects_dpo_missing_chosen(self, tmp_path): + formatter = PreferenceFormatter(output_dir=tmp_path, format_type="dpo") + assert not formatter.validate([{"prompt": "p", "rejected": "r"}]) + assert not formatter.validate([{"chosen": "c", "rejected": "r"}]) + + def test_validate_rejects_non_list(self, tmp_path): + formatter = PreferenceFormatter(output_dir=tmp_path, format_type="dpo") + assert not formatter.validate({"prompt": "p"}) + assert not formatter.validate("not-a-list") + + def test_validate_scored_format(self, tmp_path): + formatter = PreferenceFormatter(output_dir=tmp_path, format_type="score") + assert formatter.validate([{"prompt": "p", "responses": [{"response": "r", "score": 1.0}]}]) + assert not formatter.validate([{"prompt": "p"}]) + assert not formatter.validate([{"prompt": "p", "responses": "not-a-list"}]) diff --git a/tests/unit/test_pruner.py b/tests/unit/test_pruner.py new file mode 100644 index 0000000..4f38202 --- /dev/null +++ b/tests/unit/test_pruner.py @@ -0,0 +1,182 @@ +""" +Unit tests for the ContentPruner. + +The pruner is the quality gate for synthesized conversations: it decides which +conversations survive to become training data, so its filtering and structural +validation behavior is pinned down here. +""" + +import pytest + +from distill_align.core.schemas import ConversationSchema, SynthesizedTurn +from distill_align.synthesis.pruner import ContentPruner + + +@pytest.fixture +def pruner() -> ContentPruner: + """A pruner with default settings.""" + return ContentPruner() + + +def _conversation(turns: list[tuple[str, str]], **kwargs) -> ConversationSchema: + """Build a conversation from (role, content) pairs.""" + return ConversationSchema( + id=kwargs.pop("id", "conv-1"), + source_chunk_id=kwargs.pop("source_chunk_id", "chunk-1"), + turns=[SynthesizedTurn(role=role, content=content) for role, content in turns], + **kwargs, + ) + + +VALID_TURNS = [ + ("system", "You are helpful."), + ("user", "What is the capital of France?"), + ("assistant", "The capital of France is Paris, which is in Europe."), +] + + +class TestPruneConversation: + """Conversation-level pruning decisions.""" + + def test_keeps_valid_conversation(self, pruner): + conv = _conversation(VALID_TURNS) + pruned = pruner.prune_conversation(conv) + assert pruned is not None + assert pruned.id == conv.id + assert pruned.source_chunk_id == conv.source_chunk_id + assert [t.role for t in pruned.turns] == ["system", "user", "assistant"] + + def test_rejects_single_turn(self, pruner): + conv = _conversation([("user", "Hello")]) + assert pruner.prune_conversation(conv) is None + + def test_strips_filler_phrases(self, pruner): + conv = _conversation( + [ + ("user", "Sure, here is my explanation."), + ("assistant", "Here you go, have a look."), + ] + ) + pruned = pruner.prune_conversation(conv) + assert pruned is not None + assert pruned.turns[0].content == "here is my explanation." + assert pruned.turns[1].content == "have a look." + + def test_skips_empty_turns_after_clean(self, pruner): + conv = _conversation( + [ + ("system", "Sure!"), + ("user", "What is a photon?"), + ("assistant", "A photon is a quantum of light."), + ] + ) + pruned = pruner.prune_conversation(conv) + assert pruned is not None + assert [t.role for t in pruned.turns] == ["user", "assistant"] + + def test_rejects_non_alternating_roles(self, pruner): + conv = _conversation( + [ + ("user", "Question one?"), + ("user", "Question two?"), + ("assistant", "An answer that is long enough to survive pruning."), + ] + ) + assert pruner.prune_conversation(conv) is None + + def test_rejects_conversation_starting_with_assistant(self, pruner): + conv = _conversation( + [ + ("assistant", "Hello, how can I help?"), + ("user", "Question one?"), + ("assistant", "An answer that is long enough to survive pruning."), + ] + ) + assert pruner.prune_conversation(conv) is None + + def test_prune_batch_filters_low_quality(self, pruner): + good = _conversation(VALID_TURNS, id="good") + bad = _conversation([("user", "Only one")], id="bad") + result = pruner.prune_batch([good, bad]) + assert [c.id for c in result] == ["good"] + + def test_preserves_reasoning_and_confidence(self, pruner): + conv = _conversation(VALID_TURNS, reasoning_trace="trace-1", confidence_score=0.9) + pruned = pruner.prune_conversation(conv) + assert pruned is not None + assert pruned.reasoning_trace == "trace-1" + assert pruned.confidence_score == 0.9 + + +class TestExtractJson: + """JSON extraction from LLM responses.""" + + def test_extract_from_code_block(self, pruner): + content = ( + 'Some prose.\n```json\n{"conversation": [{"role": "user", "content": "Q"}]}\n```\nMore prose.' + ) + parsed = pruner.extract_json_from_response(content) + assert parsed is not None + assert parsed["conversation"][0]["role"] == "user" + + def test_extract_plain_object(self, pruner): + parsed = pruner.extract_json_from_response('Response: {"key": "value"}') + assert parsed == {"key": "value"} + + def test_extract_entire_content(self, pruner): + parsed = pruner.extract_json_from_response('{"a": 1}') + assert parsed == {"a": 1} + + def test_no_json_returns_none(self, pruner): + assert pruner.extract_json_from_response("No structured data here") is None + + +class TestValidateQuality: + """Quality scoring should flag common failure modes.""" + + def test_high_quality_conversation(self, pruner): + conv = _conversation(VALID_TURNS) + is_valid, score, issues = pruner.validate_conversation_quality(conv) + assert is_valid + assert score == 1.0 + assert issues == [] + + def test_too_few_turns(self, pruner): + conv = _conversation([("user", "Hello")]) + is_valid, score, issues = pruner.validate_conversation_quality(conv) + assert not is_valid + assert "Too few turns" in issues + + def test_too_much_filler(self, pruner): + conv = _conversation( + [ + ("user", "Sure!"), + ("assistant", "Of course!"), + ("user", "Absolutely!"), + ("assistant", "Sure."), + ] + ) + is_valid, score, issues = pruner.validate_conversation_quality(conv) + assert any("Too much filler" in issue for issue in issues) + + def test_short_content(self, pruner): + conv = _conversation([("user", "Hi"), ("assistant", "Yo")]) + is_valid, score, issues = pruner.validate_conversation_quality(conv) + # Short content alone only reduces the score; it does not reject. + assert issues == ["Content too short"] + assert score == pytest.approx(0.7) + assert is_valid + + def test_invalid_structure(self, pruner): + conv = _conversation( + [ + ("user", "Question one with enough content?"), + ("user", "Question two with enough content?"), + ("assistant", "An answer that is long enough to survive validation."), + ] + ) + is_valid, score, issues = pruner.validate_conversation_quality(conv) + # Structural issues alone only reduce the score; they do not reject. + assert issues == ["Invalid conversation structure"] + assert score == pytest.approx(0.6) + assert is_valid From d0fbbc19df8594aa4e49bdda6cdd6bc388887f84 Mon Sep 17 00:00:00 2001 From: "usehoplite[bot]" <288093033+usehoplite[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 21:08:45 +0000 Subject: [PATCH 2/2] ci: fix ruff format on new tests and add changelog entry Runs ruff format (0.9.10, matching the CI-pinned version) on the four new test files and adds an [Unreleased] CHANGELOG section covering the new tests and the PII filter import fix, satisfying the changelog-enforcer. Co-authored-by: Omar --- CHANGELOG.md | 8 ++++++++ tests/unit/test_conversation_builder.py | 4 +++- tests/unit/test_judge.py | 4 +++- tests/unit/test_pii_filter.py | 10 ++-------- tests/unit/test_pruner.py | 4 +--- 5 files changed, 17 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ac7203f..186ddee 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,14 @@ All notable changes to Distill-Align will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). +## [Unreleased] + +### Added +- Unit tests for the PII filter, pruner, judge, conversation builder, ingestion pipeline, and preference formatter (coverage 43% → 50%) + +### Fixed +- PII filter import crash caused by an inline `(?i)` regex flag in the middle of the bearer-token pattern, which broke the `scan_pii` ingestion path + ## [0.1.1] - 2026-06-18 ### Fixed diff --git a/tests/unit/test_conversation_builder.py b/tests/unit/test_conversation_builder.py index ff16675..06a395f 100644 --- a/tests/unit/test_conversation_builder.py +++ b/tests/unit/test_conversation_builder.py @@ -37,7 +37,9 @@ async def chat( raise self.error return LLMResponse(content=self.response, model=self.model) - async def complete(self, prompt: str, temperature: float = 0.7, max_tokens: int | None = None, **kwargs) -> LLMResponse: + async def complete( + self, prompt: str, temperature: float = 0.7, max_tokens: int | None = None, **kwargs + ) -> LLMResponse: return LLMResponse(content="", model=self.model) diff --git a/tests/unit/test_judge.py b/tests/unit/test_judge.py index 83630a8..449c03b 100644 --- a/tests/unit/test_judge.py +++ b/tests/unit/test_judge.py @@ -45,7 +45,9 @@ async def chat( raise LLMClientError("transient failure") return LLMResponse(content=self.response_content, model=self.model) - async def complete(self, prompt: str, temperature: float = 0.7, max_tokens: int | None = None, **kwargs) -> LLMResponse: + async def complete( + self, prompt: str, temperature: float = 0.7, max_tokens: int | None = None, **kwargs + ) -> LLMResponse: return LLMResponse(content="", model=self.model) diff --git a/tests/unit/test_pii_filter.py b/tests/unit/test_pii_filter.py index 2b23eb5..d6935ee 100644 --- a/tests/unit/test_pii_filter.py +++ b/tests/unit/test_pii_filter.py @@ -76,9 +76,7 @@ def test_aws_access_key(self, pii_filter): assert result.findings[0].severity == "critical" def test_aws_secret_key(self, pii_filter): - result = pii_filter.scan_text( - "aws_secret_access_key = wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" - ) + result = pii_filter.scan_text("aws_secret_access_key = wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY") assert _finding_types(result) == ["aws_secret_key"] def test_github_token(self, pii_filter): @@ -113,11 +111,7 @@ def test_jwt_token(self, pii_filter): assert result.findings[0].severity == "high" def test_ssh_private_key(self, pii_filter): - key = ( - "-----BEGIN RSA PRIVATE KEY-----\n" - "MIIEpAIBAAKCAQEA7eZexampleonlycontent\n" - "-----END RSA PRIVATE KEY-----" - ) + key = "-----BEGIN RSA PRIVATE KEY-----\nMIIEpAIBAAKCAQEA7eZexampleonlycontent\n-----END RSA PRIVATE KEY-----" result = pii_filter.scan_text(key) assert _finding_types(result) == ["ssh_private_key"] assert result.findings[0].severity == "critical" diff --git a/tests/unit/test_pruner.py b/tests/unit/test_pruner.py index 4f38202..f0941c0 100644 --- a/tests/unit/test_pruner.py +++ b/tests/unit/test_pruner.py @@ -112,9 +112,7 @@ class TestExtractJson: """JSON extraction from LLM responses.""" def test_extract_from_code_block(self, pruner): - content = ( - 'Some prose.\n```json\n{"conversation": [{"role": "user", "content": "Q"}]}\n```\nMore prose.' - ) + content = 'Some prose.\n```json\n{"conversation": [{"role": "user", "content": "Q"}]}\n```\nMore prose.' parsed = pruner.extract_json_from_response(content) assert parsed is not None assert parsed["conversation"][0]["role"] == "user"