|
| 1 | +"""Tests for Module A.2 chunk merge profiles.""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +import unittest |
| 6 | +from datetime import datetime, timezone |
| 7 | + |
| 8 | +from application.utils.harvester.chunk_merger import merge_chunks, requirement_ids |
| 9 | +from application.utils.harvester.models import ( |
| 10 | + ChunkInfo, |
| 11 | + Document, |
| 12 | + HeadingNode, |
| 13 | + Locator, |
| 14 | + SourceInfo, |
| 15 | +) |
| 16 | +from application.utils.harvester.schemas import ChunkingConfig |
| 17 | + |
| 18 | + |
| 19 | +def _doc(text: str, headings: list[HeadingNode] | None = None) -> Document: |
| 20 | + return Document( |
| 21 | + schema_version="0.2.0", |
| 22 | + artifact_id="art:OWASP/Test:x.md", |
| 23 | + pipeline_run_id="test", |
| 24 | + text=text, |
| 25 | + source=SourceInfo( |
| 26 | + type="github", |
| 27 | + repository="OWASP/Test", |
| 28 | + commit_sha="abc", |
| 29 | + committed_at=datetime.now(timezone.utc), |
| 30 | + ), |
| 31 | + locator=Locator(kind="repo_path", id="x.md", path="x.md"), |
| 32 | + heading_structure=headings or [], |
| 33 | + ) |
| 34 | + |
| 35 | + |
| 36 | +class ChunkMergerTests(unittest.TestCase): |
| 37 | + def test_requirement_ids_extracted(self) -> None: |
| 38 | + self.assertEqual( |
| 39 | + requirement_ids("**2.1.1** Verify passwords"), frozenset({"2.1.1"}) |
| 40 | + ) |
| 41 | + self.assertIn("2.1.1", requirement_ids("V2.1.1 something")) |
| 42 | + self.assertEqual( |
| 43 | + requirement_ids("AC-2 Account Management"), frozenset({"ac-2"}) |
| 44 | + ) |
| 45 | + self.assertIn("a.5.1", requirement_ids("ISO A.5.1 policies")) |
| 46 | + self.assertIn("3.4", requirement_ids("Req 3.4 Protect stored")) |
| 47 | + self.assertEqual( |
| 48 | + requirement_ids("Requirement 3.4.1 hashing"), frozenset({"3.4.1"}) |
| 49 | + ) |
| 50 | + |
| 51 | + def test_requirements_does_not_merge_nist_or_iso_ids(self) -> None: |
| 52 | + t1 = ("AC-2 Account Management. " * 20).strip() |
| 53 | + t2 = ("AC-3 Access Enforcement. " * 20).strip() |
| 54 | + text = t1 + "\n\n" + t2 |
| 55 | + mid = len(t1) + 2 |
| 56 | + chunks = [ |
| 57 | + ChunkInfo(text=t1, start_char_idx=0, end_char_idx=len(t1)), |
| 58 | + ChunkInfo(text=t2, start_char_idx=mid, end_char_idx=len(text)), |
| 59 | + ] |
| 60 | + headings = [ |
| 61 | + HeadingNode(level=2, text="Access Control", start_line=1, end_line=99) |
| 62 | + ] |
| 63 | + doc = _doc(text, headings) |
| 64 | + cfg = ChunkingConfig( |
| 65 | + strategy="docling", |
| 66 | + max_tokens=1200, |
| 67 | + overlap_tokens=10, |
| 68 | + merge_profile="requirements", |
| 69 | + ) |
| 70 | + out = merge_chunks(doc, chunks, cfg) |
| 71 | + self.assertEqual(len(out), 2) |
| 72 | + |
| 73 | + def test_none_profile_is_noop(self) -> None: |
| 74 | + text = "aaaa " * 50 + "\n\n" + "bbbb " * 50 |
| 75 | + mid = text.index("bbbb") |
| 76 | + chunks = [ |
| 77 | + ChunkInfo(text=text[:mid], start_char_idx=0, end_char_idx=mid), |
| 78 | + ChunkInfo(text=text[mid:], start_char_idx=mid, end_char_idx=len(text)), |
| 79 | + ] |
| 80 | + cfg = ChunkingConfig( |
| 81 | + strategy="docling", max_tokens=1200, overlap_tokens=10, merge_profile="none" |
| 82 | + ) |
| 83 | + doc = _doc(text) |
| 84 | + out = merge_chunks(doc, chunks, cfg) |
| 85 | + self.assertEqual(len(out), 2) |
| 86 | + |
| 87 | + def test_requirements_does_not_merge_across_ids(self) -> None: |
| 88 | + t1 = ("Verify 2.1.1 passwords are long. " * 20).strip() |
| 89 | + t2 = ("Verify 2.1.2 passwords allow unicode. " * 20).strip() |
| 90 | + text = t1 + "\n\n" + t2 |
| 91 | + mid = len(t1) + 2 |
| 92 | + chunks = [ |
| 93 | + ChunkInfo(text=t1, start_char_idx=0, end_char_idx=len(t1)), |
| 94 | + ChunkInfo(text=t2, start_char_idx=mid, end_char_idx=len(text)), |
| 95 | + ] |
| 96 | + headings = [ |
| 97 | + HeadingNode( |
| 98 | + level=2, text="V2.1 Password Security", start_line=1, end_line=99 |
| 99 | + ) |
| 100 | + ] |
| 101 | + doc = _doc(text, headings) |
| 102 | + cfg = ChunkingConfig( |
| 103 | + strategy="docling", |
| 104 | + max_tokens=1200, |
| 105 | + overlap_tokens=10, |
| 106 | + merge_profile="requirements", |
| 107 | + ) |
| 108 | + out = merge_chunks(doc, chunks, cfg) |
| 109 | + self.assertEqual(len(out), 2) |
| 110 | + |
| 111 | + def test_narrative_merges_same_heading_crumbs(self) -> None: |
| 112 | + # Same heading, tiny crumbs → one merged chunk under narrative profile. |
| 113 | + a = "Short intro." |
| 114 | + b = "More detail on the same topic." |
| 115 | + text = a + "\n\n" + b |
| 116 | + mid = len(a) + 2 |
| 117 | + chunks = [ |
| 118 | + ChunkInfo(text=a, start_char_idx=0, end_char_idx=len(a)), |
| 119 | + ChunkInfo(text=b, start_char_idx=mid, end_char_idx=len(text)), |
| 120 | + ] |
| 121 | + headings = [ |
| 122 | + HeadingNode(level=2, text="Password Storage", start_line=1, end_line=99) |
| 123 | + ] |
| 124 | + doc = _doc(text, headings) |
| 125 | + cfg = ChunkingConfig( |
| 126 | + strategy="docling", |
| 127 | + max_tokens=1000, |
| 128 | + overlap_tokens=10, |
| 129 | + merge_profile="narrative", |
| 130 | + ) |
| 131 | + out = merge_chunks(doc, chunks, cfg) |
| 132 | + self.assertEqual(len(out), 1) |
| 133 | + self.assertIn("Short intro", out[0].text) |
| 134 | + self.assertIn("More detail", out[0].text) |
| 135 | + |
| 136 | + def test_narrative_merges_h3_under_same_h2(self) -> None: |
| 137 | + text = ( |
| 138 | + "## Auth\n\n" |
| 139 | + "### Intro\n\n" |
| 140 | + "First crumb about auth.\n\n" |
| 141 | + "### Detail\n\n" |
| 142 | + "Second crumb still under Auth." |
| 143 | + ) |
| 144 | + # Approximate char ranges for two leaves under different ### but same ##. |
| 145 | + first = "First crumb about auth." |
| 146 | + second = "Second crumb still under Auth." |
| 147 | + i1 = text.index(first) |
| 148 | + i2 = text.index(second) |
| 149 | + chunks = [ |
| 150 | + ChunkInfo(text=first, start_char_idx=i1, end_char_idx=i1 + len(first)), |
| 151 | + ChunkInfo(text=second, start_char_idx=i2, end_char_idx=i2 + len(second)), |
| 152 | + ] |
| 153 | + headings = [ |
| 154 | + HeadingNode(level=2, text="Auth", start_line=1, end_line=20), |
| 155 | + HeadingNode(level=3, text="Intro", start_line=3, end_line=6), |
| 156 | + HeadingNode(level=3, text="Detail", start_line=8, end_line=20), |
| 157 | + ] |
| 158 | + doc = _doc(text, headings) |
| 159 | + cfg = ChunkingConfig( |
| 160 | + strategy="docling", |
| 161 | + max_tokens=1000, |
| 162 | + overlap_tokens=10, |
| 163 | + merge_profile="narrative", |
| 164 | + ) |
| 165 | + out = merge_chunks(doc, chunks, cfg) |
| 166 | + self.assertEqual(len(out), 1) |
| 167 | + |
| 168 | + def test_profile_defaults_on_config(self) -> None: |
| 169 | + req = ChunkingConfig( |
| 170 | + strategy="fixed_size", |
| 171 | + max_tokens=500, |
| 172 | + overlap_tokens=10, |
| 173 | + merge_profile="requirements", |
| 174 | + ) |
| 175 | + self.assertTrue(req.split_on_requirement_id) |
| 176 | + self.assertEqual(req.merge_min_tokens, 100) |
| 177 | + nar = ChunkingConfig( |
| 178 | + strategy="fixed_size", |
| 179 | + max_tokens=500, |
| 180 | + overlap_tokens=10, |
| 181 | + merge_profile="narrative", |
| 182 | + ) |
| 183 | + self.assertFalse(nar.split_on_requirement_id) |
| 184 | + self.assertEqual(nar.merge_min_tokens, 400) |
| 185 | + |
| 186 | + |
| 187 | +if __name__ == "__main__": |
| 188 | + unittest.main() |
0 commit comments