File size: 15,342 Bytes
ee933ab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
"""FlakeForge Integration Tests."""

import math
import os
import tempfile
from pathlib import Path
from typing import Any, Dict

import pytest

from models import (
    FlakeForgeAction,
    FlakeForgeObservation,
    FlakeForgeState,
    RunRecord,
    PatchRecord,
    RewardBreakdown,
    ROOT_CAUSE_TYPES,
    failure_mode_entropy,
)
from agent.unified_agent import (
    extract_think,
    extract_patch,
    extract_category_from_think,
    extract_confidence_from_think,
    infer_category_from_patch,
    build_unified_prompt,
)
from server.patch_applier import parse_search_replace_hunks, apply_search_replace_patch
from server.reward import (
    compute_format_reward,
    compute_compile_reward,
    compute_stability_reward,
    compute_causal_proximity_reward,
    compute_entropy_reward,
    compute_anti_hack_penalty,
    compute_reasoning_consistency,
)


class TestThinkPatchParsing:
    """Test extraction of <think> and <patch> blocks from model output."""

    def test_extract_both_blocks(self):
        response = """<think>

Root Cause: async_wait (confidence: 0.85)

Evidence: TimeoutError at test_flaky.py:12

Strategy: Increase timeout

</think>



<patch>

--- tests/test_flaky.py

<<<<<<< SEARCH

    result = await asyncio.wait_for(fetch(), timeout=0.05)

=======

    result = await asyncio.wait_for(fetch(), timeout=0.5)

>>>>>>> REPLACE

</patch>"""
        assert "Root Cause: async_wait" in extract_think(response)
        assert "<<<<<<< SEARCH" in extract_patch(response)

    def test_extract_single_json_response(self):
        response = """{

  "think": {

    "claims": [

      {

        "claim_id": "c1",

        "category": "async_wait",

        "entity": "fetch",

        "location": "tests/test_flaky.py::test_fetch",

        "ast_node_type": "Call",

        "polarity": "present",

        "predicted_effect": "The test should stop timing out.",

        "reason": "The wait_for timeout is too aggressive."

      }

    ],

    "confidence": 0.85

  },

  "patch": {

    "hunks": [

      {

        "hunk_id": "h1",

        "file": "tests/test_flaky.py",

        "search": "    result = await asyncio.wait_for(fetch(), timeout=0.05)",

        "replace": "    result = await asyncio.wait_for(fetch(), timeout=0.5)",

        "rationale": "Use a realistic timeout.",

        "addresses_claim": "c1"

      }

    ]

  }

}"""
        assert '"claims"' in extract_think(response)
        patch = extract_patch(response)
        assert "--- tests/test_flaky.py" in patch
        assert "<<<<<<< SEARCH" in patch
        assert ">>>>>>> REPLACE" in patch

    def test_extract_minimal_json_response(self):
        """7B-model friendly: only required fields, no optional keys."""
        response = """{

  "think": {

    "claims": [

      {

        "category": "async_wait",

        "entity": "fetch",

        "location": "tests/test_flaky.py::test_fetch",

        "polarity": "present",

        "reason": "Timeout too aggressive."

      }

    ],

    "confidence": 0.85

  },

  "patch": {

    "hunks": [

      {

        "file": "tests/test_flaky.py",

        "search": "    result = await asyncio.wait_for(fetch(), timeout=0.05)",

        "replace": "    result = await asyncio.wait_for(fetch(), timeout=0.5)"

      }

    ]

  }

}"""
        assert '"claims"' in extract_think(response)
        patch = extract_patch(response)
        assert "--- tests/test_flaky.py" in patch
        assert "<<<<<<< SEARCH" in patch
        assert ">>>>>>> REPLACE" in patch

    def test_extract_patch_normalizes_escaped_hunk_strings(self):
        response = """{

  "think": {

    "claims": [

      {

        "category": "concurrency",

        "entity": "submit",

        "location": "source.py::WorkerPool.submit",

        "polarity": "present",

        "reason": "Lock branch malformed"

      }

    ],

    "confidence": 0.85

  },

  "patch": {

    "hunks": [

      {

        "file": "source.py",

        "search": "        with self._lock:\\\\\\"",

        "replace": "        with self._lock:",

        "rationale": "Normalize malformed line"

      }

    ]

  }

}"""
        patch = extract_patch(response)
        assert "--- source.py" in patch
        assert "with self._lock:" in patch
        assert '\\"' not in patch

    def test_extract_empty_response(self):
        assert extract_think("") == ""
        assert extract_patch("") == ""

    def test_extract_malformed_blocks(self):
        response = "No XML blocks here"
        assert extract_think(response) == ""
        assert extract_patch(response) == ""

    def test_category_extraction(self):
        think = "Root Cause: async_wait (confidence: 0.85)"
        assert extract_category_from_think(think) == "async_wait"

    def test_category_normalization(self):
        assert extract_category_from_think("Root Cause: timeout") == "async_wait"
        assert extract_category_from_think("Root Cause: race") == "concurrency"
        assert extract_category_from_think("Root Cause: mock") == "mock_residue"
        assert extract_category_from_think("Root Cause: fixture") == "fixture_scope_leak"

    def test_confidence_extraction(self):
        think = "Root Cause: async_wait (confidence: 0.85)"
        assert extract_confidence_from_think(think) == 0.85

    def test_confidence_clamping(self):
        assert extract_confidence_from_think("confidence: 1.5") == 1.0
        assert extract_confidence_from_think("confidence: -0.5") == 0.0


class TestPatchParsing:
    """Test search/replace hunk parsing."""

    def test_single_hunk(self):
        patch = """--- tests/test_flaky.py

<<<<<<< SEARCH

    result = await asyncio.wait_for(fetch(), timeout=0.05)

=======

    result = await asyncio.wait_for(fetch(), timeout=0.5)

>>>>>>> REPLACE"""
        hunks = parse_search_replace_hunks(patch)
        assert len(hunks) == 1
        assert hunks[0].file_path == "tests/test_flaky.py"
        assert "timeout=0.05" in hunks[0].search_text
        assert "timeout=0.5" in hunks[0].replace_text

    def test_multiple_hunks(self):
        patch = """--- file1.py

<<<<<<< SEARCH

old1

=======

new1

>>>>>>> REPLACE



--- file2.py

<<<<<<< SEARCH

old2

=======

new2

>>>>>>> REPLACE"""
        hunks = parse_search_replace_hunks(patch)
        assert len(hunks) == 2
        assert hunks[0].file_path == "file1.py"
        assert hunks[1].file_path == "file2.py"

    def test_empty_patch(self):
        assert parse_search_replace_hunks("") == []
        assert parse_search_replace_hunks("   ") == []

    def test_apply_patch_to_file(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            # Create a test file
            test_file = Path(tmpdir) / "test_example.py"
            test_file.write_text("x = 1\ny = 2\nz = 3\n")

            patch = """--- test_example.py

<<<<<<< SEARCH

y = 2

=======

y = 42

>>>>>>> REPLACE"""
            result = apply_search_replace_patch(Path(tmpdir), patch)
            assert result["success"] is True
            assert result["lines_changed"] >= 1

            content = test_file.read_text()
            assert "y = 42" in content
            assert "y = 2" not in content


class TestRewardSignals:
    """Test individual reward signal computation."""

    def test_format_reward_perfect(self):
        action = FlakeForgeAction(
            raw_response="...",
            think_text="Root Cause: async_wait (confidence: 0.85)\nEvidence: ...",
            patch_text="<<<<<<< SEARCH\nold\n=======\nnew\n>>>>>>> REPLACE",
        )
        reward = compute_format_reward(action)
        assert reward == 1.0

    def test_format_reward_missing_think(self):
        action = FlakeForgeAction(
            raw_response="...",
            think_text="",
            patch_text="<<<<<<< SEARCH\nold\n=======\nnew\n>>>>>>> REPLACE",
        )
        reward = compute_format_reward(action)
        assert reward == 0.5  # Only patch block present

    def test_format_reward_empty(self):
        action = FlakeForgeAction(raw_response="garbage")
        assert compute_format_reward(action) == 0.0

    def test_compile_reward(self):
        assert compute_compile_reward(True) == 1.0
        assert compute_compile_reward(False) == -1.0
        assert compute_compile_reward(True, "SyntaxError") == -0.5

    def test_stability_full_fix(self):
        assert compute_stability_reward(0.5, 1.0) == 2.0

    def test_stability_improvement(self):
        reward = compute_stability_reward(0.3, 0.6)
        assert reward > 0

    def test_stability_regression(self):
        reward = compute_stability_reward(0.5, 0.2)
        assert reward < 0

    def test_stability_no_change(self):
        assert compute_stability_reward(0.5, 0.5) == -0.1

    def test_causal_proximity_same_file(self):
        reward = compute_causal_proximity_reward(
            patch_files=["tests/test_flaky.py"],
            failure_frontier="tests/test_flaky.py:12:test_func",
            call_chain=["test_func"],
            boundary_crossings=[],
        )
        assert reward == 1.0

    def test_causal_proximity_distant_file(self):
        reward = compute_causal_proximity_reward(
            patch_files=["utils/helpers.py"],
            failure_frontier="tests/test_flaky.py:12:test_func",
            call_chain=["test_func"],
            boundary_crossings=[],
        )
        assert reward < 0

    def test_entropy_reduction(self):
        reward = compute_entropy_reward(pre_entropy=0.8, post_entropy=0.2)
        assert reward > 0

    def test_entropy_increase(self):
        reward = compute_entropy_reward(pre_entropy=0.2, post_entropy=0.8)
        assert reward < 0

    def test_anti_hack_sleep_injection(self):
        patch = """<<<<<<< SEARCH

result = func()

=======

import time

time.sleep(1)

result = func()

>>>>>>> REPLACE"""
        penalty = compute_anti_hack_penalty(patch, ["test.py"], 3)
        assert penalty < 0

    def test_anti_hack_skip_decorator(self):
        patch = """<<<<<<< SEARCH

def test_flaky():

=======

@pytest.mark.skip

def test_flaky():

>>>>>>> REPLACE"""
        penalty = compute_anti_hack_penalty(patch, ["test.py"], 2)
        assert penalty <= -1.0

    def test_anti_hack_clean_patch(self):
        patch = """<<<<<<< SEARCH

timeout=0.05

=======

timeout=0.5

>>>>>>> REPLACE"""
        penalty = compute_anti_hack_penalty(patch, ["test.py"], 1)
        assert penalty == 0.0


class TestReasoningConsistency:
    """Test reasoning-patch consistency verification."""

    def test_exact_match(self):
        score = compute_reasoning_consistency(
            "async_wait", "async_wait", "root cause analysis", "timeout fix"
        )
        assert score == 0.5

    def test_related_match(self):
        score = compute_reasoning_consistency(
            "concurrency", "async_wait", "root cause analysis", "timeout fix"
        )
        assert score == 0.25

    def test_mismatch(self):
        score = compute_reasoning_consistency(
            "network", "fixture_scope_leak", "network issue", "fixture fix"
        )
        assert score == -0.5


class TestCategoryInference:
    """Test inferring root cause category from what the patch modifies."""

    def test_infer_async_wait(self):
        assert infer_category_from_patch("timeout=0.5") == "async_wait"

    def test_infer_concurrency(self):
        assert infer_category_from_patch("asyncio.Lock()") == "concurrency"

    def test_infer_fixture(self):
        assert infer_category_from_patch("yield fixture_data\n    teardown()") == "fixture_scope_leak"

    def test_infer_mock(self):
        assert infer_category_from_patch("monkeypatch.setattr") == "mock_residue"

    def test_infer_cache(self):
        assert infer_category_from_patch("cache_clear()") == "module_cache_pollution"


class TestFailureModeEntropy:
    """Test Shannon entropy calculation."""

    def test_single_error_type(self):
        runs = [
            RunRecord(passed=False, duration_ms=100, error_type="TimeoutError"),
            RunRecord(passed=False, duration_ms=100, error_type="TimeoutError"),
        ]
        assert failure_mode_entropy(runs) == 0.0

    def test_uniform_error_distribution(self):
        runs = [
            RunRecord(passed=False, duration_ms=100, error_type="TimeoutError"),
            RunRecord(passed=False, duration_ms=100, error_type="AssertionError"),
        ]
        assert failure_mode_entropy(runs) == 1.0

    def test_no_failures(self):
        runs = [
            RunRecord(passed=True, duration_ms=100),
            RunRecord(passed=True, duration_ms=100),
        ]
        assert failure_mode_entropy(runs) == 0.0


class TestObservationBuilding:
    """Test observation construction."""

    def test_observation_has_deep_signals(self):
        obs = FlakeForgeObservation(
            episode_id="test-001",
            test_identifier="tests/test_flaky.py::test_func",
            step=0,
            steps_remaining=8,
            test_function_source="def test_func(): pass",
            source_under_test="def func(): pass",
            module_cache_violations=["utils.py: @lru_cache on fetch"],
            async_contamination_alive=True,
        )
        assert len(obs.module_cache_violations) == 1
        assert obs.async_contamination_alive is True
        assert obs.step == 0

    def test_unified_prompt_includes_deep_signals(self):
        obs = FlakeForgeObservation(
            episode_id="test-001",
            test_identifier="tests/test_flaky.py::test_func",
            step=0,
            steps_remaining=8,
            test_function_source="def test_func(): pass",
            source_under_test="def func(): pass",
            module_cache_violations=["utils.py: @lru_cache on compute"],
            mock_residue_sites=["test_flaky.py:10 — patch() without context"],
        )
        prompt = build_unified_prompt(obs)
        assert "Module cache violations" in prompt
        assert "Mock residue sites" in prompt

    def test_observation_run_history_limit(self):
        runs = [RunRecord(passed=True, duration_ms=100) for _ in range(30)]
        obs = FlakeForgeObservation(
            episode_id="test-001",
            test_identifier="test",
            step=0,
            steps_remaining=8,
            test_function_source="",
            source_under_test="",
            run_history=runs,
        )
        assert len(obs.run_history) == 20  # Capped at 20


class TestRootCauseTypes:
    """Verify root cause taxonomy is complete."""

    def test_all_categories_present(self):
        assert "async_wait" in ROOT_CAUSE_TYPES
        assert "module_cache_pollution" in ROOT_CAUSE_TYPES
        assert "fixture_scope_leak" in ROOT_CAUSE_TYPES
        assert "mock_residue" in ROOT_CAUSE_TYPES
        assert "import_side_effect" in ROOT_CAUSE_TYPES
        assert "unknown" in ROOT_CAUSE_TYPES

    def test_category_count(self):
        assert len(ROOT_CAUSE_TYPES) == 13