CG-10520: Copy over test for runner module

pyproject.toml

@@ -112,6 +112,7 @@ dev-dependencies = [
   #    "scalene>=1.5.45",
   "filelock<4.0.0,>=3.15.4",
   "pytest>=8.3.3",
+  "pytest-asyncio>=0.23.5",
   "pytest-cov>=6.0.0,<6.0.1",
   "ruff>=0.6.8",
   "mypy[mypyc,faster-cache]>=1.13.0",

tests/unit/codegen/runner/sandbox/conftest.py

@@ -0,0 +1,58 @@
+from collections.abc import Generator
+from unittest.mock import patch
+
+import pytest
+
+from codegen.git.repo_operator.local_repo_operator import LocalRepoOperator
+from codegen.git.schemas.repo_config import RepoConfig
+from codegen.runner.models.configs import RunnerFeatureFlags
+from codegen.runner.sandbox.executor import SandboxExecutor
+from codegen.runner.sandbox.runner import SandboxRunner
+from codegen.sdk.codebase.config import CodebaseConfig, GSFeatureFlags, ProjectConfig
+from codegen.sdk.core.codebase import Codebase
+from codegen.sdk.enums import ProgrammingLanguage
+from codegen.sdk.secrets import Secrets
+
+
+@pytest.fixture
+def codebase(tmpdir, request) -> Codebase:
+    repo_id = getattr(request, "param", 1)
+    repo_config = RepoConfig(id=repo_id, name="test-repo", full_name="test-org/test-repo", organization_id=1, organization_name="test-org")
+    op = LocalRepoOperator.create_from_files(repo_path=tmpdir, files={"test.py": "a = 1"}, bot_commit=True, repo_config=repo_config)
+    projects = [ProjectConfig(repo_operator=op, programming_language=ProgrammingLanguage.PYTHON)]
+    codebase = Codebase(projects=projects)
+    return codebase
+
+
+@pytest.fixture
+def executor(codebase: Codebase) -> Generator[SandboxExecutor]:
+    with patch("codegen.runner.sandbox.executor.get_runner_feature_flags") as mock_ff:
+        mock_ff.return_value = RunnerFeatureFlags(syntax_highlight=False)
+
+        yield SandboxExecutor(codebase)
+
+
+@pytest.fixture
+def runner(codebase: Codebase, tmpdir):
+    with patch("codegen.runner.sandbox.runner.RemoteRepoOperator") as mock_op:
+        with patch.object(SandboxRunner, "_build_graph") as mock_init_codebase:
+            mock_init_codebase.return_value = codebase
+            mock_op.return_value = codebase.op
+
+            yield SandboxRunner(container_id="ta-123", repo_config=codebase.op.repo_config)
+
+
+@pytest.fixture(autouse=True)
+def mock_runner_flags():
+    with patch("codegen.runner.sandbox.executor.get_runner_feature_flags") as mock_ff:
+        mock_ff.return_value = RunnerFeatureFlags(syntax_highlight=False)
+        yield mock_ff
+
+
+@pytest.fixture(autouse=True)
+def mock_codebase_config():
+    with patch("codegen.runner.sandbox.runner.get_codebase_config") as mock_config:
+        gs_ffs = GSFeatureFlags(**RunnerFeatureFlags().model_dump())
+        secrets = Secrets(openai_key="test-key")
+        mock_config.return_value = CodebaseConfig(secrets=secrets, feature_flags=gs_ffs)
+        yield mock_config

tests/unit/codegen/runner/sandbox/test_executor.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from codegen.git.models.codemod_context import CodemodContext
+from codegen.runner.models.codemod import GroupingConfig
+from codegen.runner.sandbox.executor import SandboxExecutor
+from codegen.sdk.codebase.config import SessionOptions
+from codegen.sdk.codebase.flagging.code_flag import CodeFlag
+from codegen.sdk.codebase.flagging.groupers.enums import GroupBy
+from codegen.shared.compilation.string_to_code import create_execute_function_from_codeblock
+
+
+@pytest.mark.asyncio
+async def test_execute_func_pass_in_codemod_context_takes_priority(executor: SandboxExecutor):
+    codemod_context = CodemodContext(
+        CODEMOD_LINK="http://codegen.sh/codemod/5678",
+    )
+    mock_source = """
+print(context.CODEMOD_LINK)
+"""
+    custom_scope = {"context": codemod_context}
+    code_to_exec = create_execute_function_from_codeblock(codeblock=mock_source, custom_scope=custom_scope)
+    mock_log = MagicMock()
+    executor.codebase.log = mock_log
+
+    result = await executor.execute(code_to_exec)
+
+    assert result is not None
+
+    assert mock_log.call_count == 1
+    assert mock_log.call_args_list[0][0][0] == "http://codegen.sh/codemod/5678"
+
+
+# @pytest.mark.asyncio
+# async def test_init_execute_func_with_pull_request_context(executor: SandboxExecutor):
+#     mock_source = """
+# print(context.PULL_REQUEST.head.ref)
+# print(context.PULL_REQUEST.base.ref)
+# """
+#     mock_cm_run = MagicMock(epic=MagicMock(id=1234, link="link", user=MagicMock(github_username="user")), codemod_version=MagicMock(source=mock_source))
+#     mock_pull = MagicMock(spec=GithubWebhookPR, head=MagicMock(ref="test-head"), base=MagicMock(ref="test-base"))
+#     codemod_context = get_codemod_context(cm_run=mock_cm_run, pull_request=mock_pull)
+#     custom_scope = {"context": codemod_context}
+#     code_to_exec = create_execute_function_from_codeblock(codeblock=mock_source, custom_scope=custom_scope)
+#     mock_log = MagicMock()
+#     executor.codebase.log = mock_log
+#
+#     result = await executor.execute(code_to_exec)
+#
+#     assert result is not None
+#     assert mock_log.call_count == 2
+#     assert mock_log.call_args_list[0][0][0] == "test-head"
+#     assert mock_log.call_args_list[1][0][0] == "test-base"
+#
+#
+# @pytest.mark.asyncio
+# async def test_init_execute_func_with_pull_request_context_mock_codebase(executor: SandboxExecutor):
+#     mock_source = """
+# print(context.PULL_REQUEST.head.ref)
+# print(context.PULL_REQUEST.base.ref)
+#     """
+#     mock_cm_run = MagicMock(epic=MagicMock(id=1234, link="link", user=MagicMock(github_username="user")), codemod_version=MagicMock(source=mock_source))
+#     mock_pull = MagicMock(spec=GithubWebhookPR, head=MagicMock(ref="test-head"), base=MagicMock(ref="test-base"))
+#     codemod_context = get_codemod_context(cm_run=mock_cm_run, pull_request=mock_pull)
+#     custom_scope = {"context": codemod_context}
+#     code_to_exec = create_execute_function_from_codeblock(codeblock=mock_source, custom_scope=custom_scope)
+#
+#     result = await executor.execute(code_to_exec)
+#
+#     # validate
+#     assert result is not None
+#     assert (
+#         result.logs
+#         == """
+# test-head
+# test-base
+# """.lstrip()
+#     )
+
+
+@pytest.mark.asyncio
+async def test_run_max_preview_time_exceeded_sets_observation_meta(executor: SandboxExecutor):
+    mock_source = """
+codebase.files[0].edit("a = 2")
+"""
+    code_to_exec = create_execute_function_from_codeblock(codeblock=mock_source)
+    result = await executor.execute(code_to_exec, session_options=SessionOptions(max_seconds=0))
+
+    assert result.is_complete
+    assert result.observation_meta == {"flags": [], "stop_codemod_exception_type": "MaxPreviewTimeExceeded", "threshold": 0}
+
+
+@pytest.mark.asyncio
+async def test_run_max_ai_requests_error_sets_observation_meta(executor: SandboxExecutor):
+    mock_source = """
+codebase.ai("tell me a joke")
+"""
+    code_to_exec = create_execute_function_from_codeblock(codeblock=mock_source)
+    result = await executor.execute(code_to_exec, session_options=SessionOptions(max_ai_requests=0))
+
+    assert result.is_complete
+    assert result.observation_meta == {"flags": [], "stop_codemod_exception_type": "MaxAIRequestsError", "threshold": 0}
+
+
+@pytest.mark.asyncio
+async def test_run_max_transactions_exceeded_sets_observation_meta(executor: SandboxExecutor):
+    mock_source = """
+codebase.files[0].edit("a = 2")
+"""
+
+    code_to_exec = create_execute_function_from_codeblock(codeblock=mock_source)
+    result = await executor.execute(code_to_exec, session_options=SessionOptions(max_transactions=0))
+
+    assert result.is_complete
+    assert result.observation_meta == {"flags": [], "stop_codemod_exception_type": "MaxTransactionsExceeded", "threshold": 0}
+
+
+@pytest.mark.asyncio
+async def test_find_flag_groups_with_subdirectories(executor: SandboxExecutor):
+    groups = await executor.find_flag_groups(
+        code_flags=[
+            CodeFlag(
+                symbol=MagicMock(file=MagicMock(filepath="subdir1/file1.py")),
+                message="message",
+            ),
+            CodeFlag(
+                symbol=MagicMock(file=MagicMock(filepath="subdir2/file1.py")),
+                message="message",
+            ),
+            CodeFlag(
+                symbol=MagicMock(file=MagicMock(filepath="subdir3/file1.py")),
+                message="message",
+            ),
+            CodeFlag(
+                symbol=MagicMock(file=MagicMock(filepath="subdir3/file2.py")),
+                message="message",
+            ),
+        ],
+        grouping_config=GroupingConfig(subdirectories=["subdir1", "subdir2"]),
+    )
+    assert len(groups) == 1
+    assert len(groups[0].flags) == 2
+    assert groups[0].flags[0].filepath == "subdir1/file1.py"
+    assert groups[0].flags[1].filepath == "subdir2/file1.py"
+
+
+@pytest.mark.asyncio
+async def test_find_flag_groups_with_group_by(executor: SandboxExecutor):
+    groups = await executor.find_flag_groups(
+        code_flags=[
+            CodeFlag(
+                symbol=MagicMock(file=MagicMock(filepath="subdir1/file1.py")),
+                message="message",
+            ),
+            CodeFlag(
+                symbol=MagicMock(file=MagicMock(filepath="subdir2/file1.py")),
+                message="message",
+            ),
+            CodeFlag(
+                symbol=MagicMock(file=MagicMock(filepath="subdir3/file1.py")),
+                message="message",
+            ),
+            CodeFlag(
+                symbol=MagicMock(file=MagicMock(filepath="subdir1/file1.py")),
+                message="message",
+            ),
+        ],
+        grouping_config=GroupingConfig(group_by=GroupBy.FILE),
+    )
+    assert len(groups) == 3
+    assert groups[0].segment == "subdir1/file1.py"
+    assert groups[1].segment == "subdir2/file1.py"
+    assert groups[2].segment == "subdir3/file1.py"
+
+    assert len(groups[0].flags) == 2
+    assert len(groups[1].flags) == 1
+    assert len(groups[2].flags) == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("codebase", [121], indirect=True)
+async def test_find_flag_groups_with_group_by_app(executor: SandboxExecutor):
+    groups = await executor.find_flag_groups(
+        code_flags=[
+            CodeFlag(
+                symbol=MagicMock(file=MagicMock(filepath="a/b/app1/test1.py")),
+                message="message",
+            ),
+            CodeFlag(
+                symbol=MagicMock(file=MagicMock(filepath="a/b/app2/test1.py")),
+                message="message",
+            ),
+            CodeFlag(
+                symbol=MagicMock(file=MagicMock(filepath="a/b/app3/test1.py")),
+                message="message",
+            ),
+            CodeFlag(
+                symbol=MagicMock(file=MagicMock(filepath="a/b/app2/test2.py")),
+                message="message",
+            ),
+        ],
+        grouping_config=GroupingConfig(group_by=GroupBy.APP),
+    )
+    count_by_segment = {group.segment: len(group.flags) for group in groups}
+    assert count_by_segment == {"a/b/app1": 1, "a/b/app2": 2, "a/b/app3": 1}
+
+
+@pytest.mark.skip(reason="TODO: add max_prs as part of find_flag_groups")
+@pytest.mark.asyncio
+async def test_find_flag_groups_with_max_prs(executor: SandboxExecutor):
+    groups = await executor.find_flag_groups(
+        code_flags=[
+            CodeFlag(
+                symbol=MagicMock(file=MagicMock(filepath="subdir1/file1.py")),
+                message="message",
+            ),
+            CodeFlag(
+                symbol=MagicMock(file=MagicMock(filepath="subdir2/file1.py")),
+                message="message",
+            ),
+        ],
+        grouping_config=GroupingConfig(group_by=GroupBy.FILE, max_prs=0),
+    )
+    assert len(groups) == 0