codeflash/tests/test_languages/test_java/test_tracer_exit_codes.py at 470482e82474f0210c9de09ef010ac2b17b0db30 · codeflash-ai/codeflash · GitHub

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
from __future__ import annotations

from typing import TYPE_CHECKING
from unittest.mock import MagicMock, patch

if TYPE_CHECKING:
    from pathlib import Path

import pytest

from codeflash.languages.java.tracer import JavaTracer, _run_java_with_graceful_timeout


class TestRunJavaWithGracefulTimeout:
    def test_returns_zero_on_success(self) -> None:
        mock_result = MagicMock()
        mock_result.returncode = 0
        with patch("codeflash.languages.java.tracer.subprocess.run", return_value=mock_result):
            rc = _run_java_with_graceful_timeout(["java", "-version"], {}, 0, "test")
        assert rc == 0

    def test_returns_nonzero_on_failure(self) -> None:
        mock_result = MagicMock()
        mock_result.returncode = 1
        with patch("codeflash.languages.java.tracer.subprocess.run", return_value=mock_result):
            rc = _run_java_with_graceful_timeout(["java", "-version"], {}, 0, "test")
        assert rc == 1

    def test_returns_exit_code_137_oom_kill(self) -> None:
        mock_result = MagicMock()
        mock_result.returncode = 137
        with patch("codeflash.languages.java.tracer.subprocess.run", return_value=mock_result):
            rc = _run_java_with_graceful_timeout(["java", "-version"], {}, 0, "test")
        assert rc == 137

    def test_timeout_path_returns_zero_on_success(self) -> None:
        mock_proc = MagicMock()
        mock_proc.returncode = 0
        with patch("codeflash.languages.java.tracer.subprocess.Popen", return_value=mock_proc):
            rc = _run_java_with_graceful_timeout(["java", "-version"], {}, 60, "test")
        assert rc == 0

    def test_timeout_path_returns_nonzero_on_failure(self) -> None:
        mock_proc = MagicMock()
        mock_proc.returncode = 1
        with patch("codeflash.languages.java.tracer.subprocess.Popen", return_value=mock_proc):
            rc = _run_java_with_graceful_timeout(["java", "-version"], {}, 60, "test")
        assert rc == 1

    def test_timeout_returns_negative_one(self) -> None:
        import subprocess

        mock_proc = MagicMock()
        # First wait() times out, SIGTERM wait succeeds
        mock_proc.wait.side_effect = [
            subprocess.TimeoutExpired(cmd="java", timeout=60),
            None,  # SIGTERM wait succeeds
        ]
        with patch("codeflash.languages.java.tracer.subprocess.Popen", return_value=mock_proc):
            rc = _run_java_with_graceful_timeout(["java", "-version"], {}, 60, "test")
        assert rc == -1

    def test_timeout_sends_sigterm_then_sigkill(self) -> None:
        import signal
        import subprocess

        mock_proc = MagicMock()
        # First wait() times out, SIGTERM wait also times out
        mock_proc.wait.side_effect = [
            subprocess.TimeoutExpired(cmd="java", timeout=60),
            subprocess.TimeoutExpired(cmd="java", timeout=5),
            None,
        ]
        with patch("codeflash.languages.java.tracer.subprocess.Popen", return_value=mock_proc):
            rc = _run_java_with_graceful_timeout(["java", "-version"], {}, 60, "test")

        assert rc == -1
        mock_proc.send_signal.assert_called_once_with(signal.SIGTERM)
        mock_proc.kill.assert_called_once()


class TestJavaTracerExitCodeHandling:
    def test_stage1_failure_continues(self, tmp_path: Path) -> None:
        trace_db_path = (tmp_path / "trace.db").resolve()
        tracer = JavaTracer()

        # Stage 1 fails (exit code 1), Stage 2 succeeds (exit code 0)
        exit_codes = iter([1, 0])

        def mock_run_timeout(java_command: list[str], env: dict, timeout: int, stage_name: str) -> int:
            rc = next(exit_codes)
            if stage_name == "Argument capture":
                trace_db_path.write_bytes(b"fake-db")
            return rc

        with (
            patch("codeflash.languages.java.tracer._run_java_with_graceful_timeout", side_effect=mock_run_timeout),
            patch.object(tracer, "build_jfr_env", return_value={}),
            patch.object(tracer, "build_agent_env", return_value={}),
            patch.object(tracer, "create_tracer_config", return_value=tmp_path / "config.json"),
        ):
            trace_db, _jfr_file = tracer.trace(
                java_command=["java", "-cp", ".", "Main"], trace_db_path=trace_db_path, packages=["com.example"]
            )
        # Should complete despite Stage 1 failure
        assert trace_db == trace_db_path

    def test_stage2_failure_raises(self, tmp_path: Path) -> None:
        trace_db_path = (tmp_path / "trace.db").resolve()
        tracer = JavaTracer()

        # Stage 1 succeeds (exit code 0), Stage 2 fails (exit code 1)
        exit_codes = iter([0, 1])

        def mock_run_timeout(java_command: list[str], env: dict, timeout: int, stage_name: str) -> int:
            return next(exit_codes)

        with (
            patch("codeflash.languages.java.tracer._run_java_with_graceful_timeout", side_effect=mock_run_timeout),
            patch.object(tracer, "build_jfr_env", return_value={}),
            patch.object(tracer, "build_agent_env", return_value={}),
            patch.object(tracer, "create_tracer_config", return_value=tmp_path / "config.json"),
            pytest.raises(RuntimeError, match="Argument capture failed with exit code 1"),
        ):
            tracer.trace(
                java_command=["java", "-cp", ".", "Main"], trace_db_path=trace_db_path, packages=["com.example"]
            )

    def test_both_stages_succeed(self, tmp_path: Path) -> None:
        trace_db_path = (tmp_path / "trace.db").resolve()
        tracer = JavaTracer()

        def mock_run_timeout(java_command: list[str], env: dict, timeout: int, stage_name: str) -> int:
            if stage_name == "Argument capture":
                trace_db_path.write_bytes(b"fake-db")
            return 0

        with (
            patch("codeflash.languages.java.tracer._run_java_with_graceful_timeout", side_effect=mock_run_timeout),
            patch.object(tracer, "build_jfr_env", return_value={}),
            patch.object(tracer, "build_agent_env", return_value={}),
            patch.object(tracer, "create_tracer_config", return_value=tmp_path / "config.json"),
        ):
            trace_db, _jfr_file = tracer.trace(
                java_command=["java", "-cp", ".", "Main"], trace_db_path=trace_db_path, packages=["com.example"]
            )
        assert trace_db == trace_db_path