update tests and restore sandbox tests

Author: bagel897

tests/integration/codegen/runner/conftest.py

@@ -0,0 +1,39 @@
+from collections.abc import Generator
+from unittest.mock import Mock
+
+import pytest
+
+from graph_sitter.git.clients.git_repo_client import GitRepoClient
+from graph_sitter.git.repo_operator.repo_operator import RepoOperator
+from graph_sitter.git.schemas.enums import SetupOption
+from graph_sitter.git.schemas.repo_config import RepoConfig
+from graph_sitter.runner.clients.codebase_client import CodebaseClient
+from graph_sitter.shared.enums.programming_language import ProgrammingLanguage
+from graph_sitter.shared.network.port import get_free_port
+
+
+@pytest.fixture()
+def repo_config(tmpdir) -> Generator[RepoConfig, None, None]:
+    yield RepoConfig(
+        name="Kevin-s-Adventure-Game",
+        full_name="codegen-sh/Kevin-s-Adventure-Game",
+        language=ProgrammingLanguage.PYTHON,
+        base_dir=str(tmpdir),
+    )
+
+
+@pytest.fixture
+def op(repo_config: RepoConfig) -> Generator[RepoOperator, None, None]:
+    yield RepoOperator(repo_config=repo_config, setup_option=SetupOption.PULL_OR_CLONE)
+
+
+@pytest.fixture
+def git_repo_client(op: RepoOperator, repo_config: RepoConfig) -> Generator[GitRepoClient, None, None]:
+    yield GitRepoClient(repo_config=repo_config, access_token=op.access_token)
+
+
+@pytest.fixture
+def codebase_client(repo_config: RepoConfig) -> Generator[CodebaseClient, None, None]:
+    sb_client = CodebaseClient(repo_config=repo_config, port=get_free_port())
+    sb_client.runner = Mock()
+    yield sb_client

tests/integration/codegen/runner/test_create_branch.py

@@ -0,0 +1,63 @@
+import uuid
+from http import HTTPStatus
+
+import pytest
+
+from graph_sitter.git.clients.git_repo_client import GitRepoClient
+from graph_sitter.git.repo_operator.repo_operator import RepoOperator
+from graph_sitter.runner.clients.codebase_client import CodebaseClient
+from graph_sitter.runner.models.apis import BRANCH_ENDPOINT, CreateBranchRequest, CreateBranchResponse
+from graph_sitter.runner.models.codemod import BranchConfig, Codemod, GroupingConfig
+
+
+@pytest.mark.asyncio
+@pytest.mark.timeout(60)
+async def test_create_branch(codebase_client: CodebaseClient, git_repo_client: GitRepoClient, op: RepoOperator):
+    # set-up
+    codemod_source = """
+for file in codebase.files:
+    new_content = "🌈" + "\\n" + file.content
+    file.edit(new_content)
+    break
+"""
+    test_branch_name = f"codegen-test-create-branch-{uuid.uuid1()}"
+    request = CreateBranchRequest(
+        codemod=Codemod(user_code=codemod_source),
+        commit_msg="Create branch test",
+        grouping_config=GroupingConfig(),
+        branch_config=BranchConfig(branch_name=test_branch_name),
+    )
+
+    # execute
+    response = codebase_client.post(endpoint=BRANCH_ENDPOINT, data=request.model_dump())
+    assert response.status_code == HTTPStatus.OK
+
+    # verify
+    result = CreateBranchResponse.model_validate(response.json())
+    assert len(result.results) == 1
+    assert result.results[0].is_complete
+    assert result.results[0].error is None
+    assert result.results[0].logs == ""
+    assert result.results[0].observation is not None
+
+    # verify changed files
+    patch = result.results[0].observation
+    lines = patch.split("\n")
+    added_lines = [line[1:] for line in lines if line.startswith("+") and len(line) > 1]
+    assert "🌈" in added_lines
+
+    # verify returned branch
+    assert len(result.branches) == 1
+    branch = result.branches[0]
+    assert branch.base_branch == "main"
+    assert branch.head_ref == test_branch_name
+
+    # verify remote branch
+    remote_branch = git_repo_client.repo.get_branch(test_branch_name)
+    assert remote_branch is not None
+    assert remote_branch.name == test_branch_name
+    assert remote_branch.commit.commit.message == "[Codegen] Create branch test"
+
+    # clean-up
+    remote = op.git_cli.remote(name="origin")
+    remote.push([f":refs/heads/{test_branch_name}"])  # The colon prefix means delete

tests/integration/codegen/runner/test_create_branch_with_grouping.py

@@ -0,0 +1,58 @@
+import uuid
+from http import HTTPStatus
+
+import pytest
+
+from graph_sitter.codebase.flagging.groupers.enums import GroupBy
+from graph_sitter.git.clients.git_repo_client import GitRepoClient
+from graph_sitter.git.repo_operator.repo_operator import RepoOperator
+from graph_sitter.runner.clients.codebase_client import CodebaseClient
+from graph_sitter.runner.models.apis import BRANCH_ENDPOINT, CreateBranchRequest, CreateBranchResponse
+from graph_sitter.runner.models.codemod import BranchConfig, Codemod, GroupingConfig
+
+
+@pytest.mark.timeout(120)
+@pytest.mark.parametrize("group_by", [GroupBy.INSTANCE, GroupBy.FILE])
+def test_create_branch_with_grouping(codebase_client: CodebaseClient, git_repo_client: GitRepoClient, op: RepoOperator, group_by: GroupBy):
+    codemod_source = """
+for file in codebase.files[:5]:
+    flag = codebase.flag_instance(file)
+    if codebase.should_fix(flag):
+        new_content = "🌈" + "\\n" + file.content
+        file.edit(new_content)
+"""
+    commit_msg = "Create branch with grouping test"
+    test_branch_name = f"codegen-{uuid.uuid1()}"
+    request = CreateBranchRequest(
+        codemod=Codemod(user_code=codemod_source),
+        commit_msg=commit_msg,
+        grouping_config=GroupingConfig(group_by=group_by),
+        branch_config=BranchConfig(branch_name=test_branch_name),
+    )
+
+    # execute
+    response = codebase_client.post(endpoint=BRANCH_ENDPOINT, data=request.model_dump())
+    assert response.status_code == HTTPStatus.OK
+
+    # verify
+    result = CreateBranchResponse.model_validate(response.json())
+    assert len(result.results) == 5
+    assert len(result.branches) == 5
+
+    for i, branch in enumerate(result.branches):
+        actual_branch_suffix = "-".join(branch.head_ref.split("-")[-2:])
+        expected_branch_suffix = f"group-{i}"
+        assert expected_branch_suffix == actual_branch_suffix
+
+        remote_branch = git_repo_client.repo.get_branch(branch.head_ref)
+        assert remote_branch is not None
+        assert remote_branch.name == branch.head_ref
+        assert remote_branch.commit.commit.message == f"[Codegen] {commit_msg}"
+        assert remote_branch.commit.commit.author.name == "codegen-sh[bot]"
+
+        comparison = git_repo_client.repo.compare(base=branch.base_branch, head=branch.head_ref)
+        assert "+🌈" in comparison.files[0].patch
+
+        # clean-up
+        remote = op.git_cli.remote(name="origin")
+        remote.push([f":refs/heads/{branch.head_ref}"])

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

@@ -0,0 +1,34 @@
+from collections.abc import Generator
+from unittest.mock import patch
+
+import pytest
+
+from graph_sitter.codebase.config import ProjectConfig
+from graph_sitter.core.codebase import Codebase
+from graph_sitter.git.repo_operator.repo_operator import RepoOperator
+from graph_sitter.runner.sandbox.executor import SandboxExecutor
+from graph_sitter.runner.sandbox.runner import SandboxRunner
+from graph_sitter.shared.enums.programming_language import ProgrammingLanguage
+
+
+@pytest.fixture
+def codebase(tmpdir) -> Codebase:
+    op = RepoOperator.create_from_files(repo_path=f"{tmpdir}/test-repo", files={"test.py": "a = 1"}, bot_commit=True)
+    projects = [ProjectConfig(repo_operator=op, programming_language=ProgrammingLanguage.PYTHON)]
+    codebase = Codebase(projects=projects)
+    return codebase
+
+
+@pytest.fixture
+def executor(codebase: Codebase) -> Generator[SandboxExecutor]:
+    yield SandboxExecutor(codebase)
+
+
+@pytest.fixture
+def runner(codebase: Codebase, tmpdir):
+    with patch("graph_sitter.runner.sandbox.runner.RepoOperator") 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(repo_config=codebase.op.repo_config)