diff --git a/packages/bot/CHANGELOG.md b/packages/bot/CHANGELOG.md index af7b5a2..d9bf113 100644 --- a/packages/bot/CHANGELOG.md +++ b/packages/bot/CHANGELOG.md @@ -2,6 +2,12 @@ ### Unreleased +### 0.3.0 + +#### Added + +- Support direct download method for code + ### 0.2.3 #### Added diff --git a/packages/bot/pyproject.toml b/packages/bot/pyproject.toml index 9d3eac6..50ed249 100644 --- a/packages/bot/pyproject.toml +++ b/packages/bot/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "automa-bot" -version = "0.2.3" +version = "0.3.0" authors = [{ name = "Sunkara, Inc.", email = "engineering@automa.app" }] description = "Bot helpers for Automa" diff --git a/packages/bot/src/automa/bot/resources/code.py b/packages/bot/src/automa/bot/resources/code.py index ed84325..bf7310d 100644 --- a/packages/bot/src/automa/bot/resources/code.py +++ b/packages/bot/src/automa/bot/resources/code.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import subprocess import tarfile from asyncio import to_thread @@ -77,6 +78,42 @@ def _read_token(self, folder: str) -> str | None: return token + def _read_base_commit(self, folder: str) -> str | None: + base_commit = None + + try: + with open( + f"{folder}/.git/automa_proposal_base_commit", "r", encoding="utf8" + ) as f: + base_commit = f.read().strip() + except FileNotFoundError: + pass + + return base_commit + + def _clone_code(self, folder: str, url: str) -> None: + path = Path(folder) + + rmtree(folder, ignore_errors=True) + path.mkdir(parents=True, exist_ok=True) + + subprocess.run( + ["git", "clone", "--depth=1", url, "."], + cwd=path, + check=True, + ) + + # Note down the base commit + result = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=path, + capture_output=True, + text=True, + check=True, + ) + + self._write_base_commit(folder, result.stdout.strip()) + def _extract_download(self, folder: str) -> None: rmtree(folder, ignore_errors=True) Path(folder).mkdir(parents=True, exist_ok=True) @@ -89,13 +126,24 @@ def _write_token(self, folder: str, token: str) -> None: with open(f"{folder}/.git/automa_proposal_token", "w", encoding="utf8") as f: f.write(token) + def _write_base_commit(self, folder: str, base_commit: str) -> None: + # Save the base commit for later use + with open( + f"{folder}/.git/automa_proposal_base_commit", "w", encoding="utf8" + ) as f: + f.write(base_commit) + class CodeResource(SyncAPIResource, BaseCodeResource): def cleanup(self, body: CodeCleanupParams) -> None: folder = self._path(body["task"]) rmtree(folder, ignore_errors=True) - remove(f"{folder}.tar.gz") + + try: + remove(f"{folder}.tar.gz") + except FileNotFoundError: + pass def download( self, body: CodeDownloadParams, *, options: RequestOptions = {} @@ -112,20 +160,31 @@ def download( "json": body, "headers": { **options.get("headers", {}), - "Accept": "application/gzip", }, }, ) as response: - token = response.headers["x-automa-proposal-token"] + token = response.headers.get("x-automa-proposal-token") + content_type = response.headers.get("Content-Type", "") + + if content_type.startswith("application/json"): + content = b"".join(response.iter_bytes()) - rmtree(path, ignore_errors=True) - makedirs(path, exist_ok=True) + json_data = json.loads(content.decode("utf-8")) - with open(archive_path, "wb") as archive: - for chunk in response.iter_bytes(chunk_size=8192): - archive.write(chunk) + self._clone_code(path, json_data["url"]) + elif content_type.startswith("application/gzip"): + rmtree(path, ignore_errors=True) + makedirs(path, exist_ok=True) - self._extract_download(path) + with open(archive_path, "wb") as archive: + for chunk in response.iter_bytes(chunk_size=8192): + archive.write(chunk) + + self._extract_download(path) + else: + raise ValueError( + f"Unexpected Content-Type: {content_type} while downloading code." + ) # Save the proposal token for later use self._write_token(path, token) @@ -135,6 +194,7 @@ def download( def propose(self, body: CodeProposeParams, *, options: RequestOptions = {}): path = self._path(body["task"]) token = self._read_token(path) + base_commit = self._read_base_commit(path) if not token: raise ValueError("Failed to read the stored proposal token") @@ -149,6 +209,7 @@ def propose(self, body: CodeProposeParams, *, options: RequestOptions = {}): **body.get("proposal", {}), "token": token, "diff": diff, + **({"base_commit": base_commit} if base_commit else {}), }, }, options=options, @@ -160,7 +221,11 @@ async def cleanup(self, body: CodeCleanupParams) -> None: folder = self._path(body["task"]) await to_thread(rmtree, folder, ignore_errors=True) - await to_thread(remove, f"{folder}.tar.gz") + + try: + await to_thread(remove, f"{folder}.tar.gz") + except FileNotFoundError: + pass async def download( self, body: CodeDownloadParams, *, options: RequestOptions = {} @@ -177,20 +242,34 @@ async def download( "json": body, "headers": { **options.get("headers", {}), - "Accept": "application/gzip", }, }, ) as response: - token = response.headers["x-automa-proposal-token"] + token = response.headers.get("x-automa-proposal-token") + content_type = response.headers.get("Content-Type", "") + + if content_type.startswith("application/json"): + content = b"" + + async for chunk in response.aiter_bytes(): + content += chunk + + json_data = json.loads(content.decode("utf-8")) - await to_thread(rmtree, path, ignore_errors=True) - await to_thread(makedirs, path, exist_ok=True) + await to_thread(self._clone_code, path, json_data["url"]) + elif content_type.startswith("application/gzip"): + await to_thread(rmtree, path, ignore_errors=True) + await to_thread(makedirs, path, exist_ok=True) - with open(archive_path, "wb") as archive: - async for chunk in response.aiter_bytes(chunk_size=8192): - await to_thread(archive.write, chunk) + with open(archive_path, "wb") as archive: + async for chunk in response.aiter_bytes(chunk_size=8192): + await to_thread(archive.write, chunk) - await to_thread(self._extract_download, path) + await to_thread(self._extract_download, path) + else: + raise ValueError( + f"Unexpected Content-Type: {content_type} while downloading code." + ) # Save the proposal token for later use await to_thread(self._write_token, path, token) @@ -200,6 +279,7 @@ async def download( async def propose(self, body: CodeProposeParams, *, options: RequestOptions = {}): path = self._path(body["task"]) token = await to_thread(self._read_token, path) + base_commit = await to_thread(self._read_base_commit, path) if not token: raise ValueError("Failed to read the stored proposal token") @@ -214,6 +294,7 @@ async def propose(self, body: CodeProposeParams, *, options: RequestOptions = {} **body.get("proposal", {}), "token": token, "diff": diff, + **({"base_commit": base_commit} if base_commit else {}), }, }, options=options, diff --git a/packages/bot/tests/fixtures/download_git/HEAD b/packages/bot/tests/fixtures/download_git/HEAD new file mode 100644 index 0000000..cb089cd --- /dev/null +++ b/packages/bot/tests/fixtures/download_git/HEAD @@ -0,0 +1 @@ +ref: refs/heads/master diff --git a/packages/bot/tests/fixtures/download_git/index b/packages/bot/tests/fixtures/download_git/index new file mode 100644 index 0000000..083dc80 Binary files /dev/null and b/packages/bot/tests/fixtures/download_git/index differ diff --git a/packages/bot/tests/fixtures/download_git/objects/cc/3f46ae7fdf71747b66b3e4272c0e5fe290d116 b/packages/bot/tests/fixtures/download_git/objects/cc/3f46ae7fdf71747b66b3e4272c0e5fe290d116 new file mode 100644 index 0000000..51a0b9f Binary files /dev/null and b/packages/bot/tests/fixtures/download_git/objects/cc/3f46ae7fdf71747b66b3e4272c0e5fe290d116 differ diff --git a/packages/bot/tests/fixtures/download_git/objects/e6/9de29bb2d1d6434b8b29ae775ad8c2e48c5391 b/packages/bot/tests/fixtures/download_git/objects/e6/9de29bb2d1d6434b8b29ae775ad8c2e48c5391 new file mode 100644 index 0000000..7112238 Binary files /dev/null and b/packages/bot/tests/fixtures/download_git/objects/e6/9de29bb2d1d6434b8b29ae775ad8c2e48c5391 differ diff --git a/packages/bot/tests/fixtures/download_git/objects/f9/3e3a1a1525fb5b91020da86e44810c87a2d7bc b/packages/bot/tests/fixtures/download_git/objects/f9/3e3a1a1525fb5b91020da86e44810c87a2d7bc new file mode 100644 index 0000000..91fccd4 Binary files /dev/null and b/packages/bot/tests/fixtures/download_git/objects/f9/3e3a1a1525fb5b91020da86e44810c87a2d7bc differ diff --git a/packages/bot/tests/fixtures/download_git/refs/heads/master b/packages/bot/tests/fixtures/download_git/refs/heads/master new file mode 100644 index 0000000..be3d3dc --- /dev/null +++ b/packages/bot/tests/fixtures/download_git/refs/heads/master @@ -0,0 +1 @@ +cc3f46ae7fdf71747b66b3e4272c0e5fe290d116 diff --git a/packages/bot/tests/resources/test_code.py b/packages/bot/tests/resources/test_code.py index eea7e0a..c921feb 100644 --- a/packages/bot/tests/resources/test_code.py +++ b/packages/bot/tests/resources/test_code.py @@ -11,6 +11,7 @@ folder = "/tmp/automa/tasks/28" proposal_token_file = f"{folder}/.git/automa_proposal_token" +proposal_base_commit_file = f"{folder}/.git/automa_proposal_base_commit" @pytest.fixture @@ -41,6 +42,14 @@ def async_code_resource(): rmtree(path.dirname(folder), ignore_errors=True) +@pytest.fixture +def fixture_git(): + tests_folder = Path(__file__).parent.parent + git_path = tests_folder / "fixtures" / "download_git" + + yield git_path + + @pytest.fixture def fixture_tarfile(): tests_folder = Path(__file__).parent.parent @@ -65,7 +74,7 @@ def fixture_tarfile(): rmtree(fixture / ".git", ignore_errors=True) -def test_cleanup(code_resource): +def test_cleanup_proxy(code_resource): makedirs(folder, exist_ok=True) with open(f"{folder}.tar.gz", "w") as f: f.write("ghijkl") @@ -80,6 +89,19 @@ def test_cleanup(code_resource): assert not path.exists(f"{folder}.tar.gz") +def test_cleanup_direct(code_resource): + makedirs(folder, exist_ok=True) + + assert path.exists(folder) + assert not path.exists(f"{folder}.tar.gz") + + # Call cleanup + code_resource.cleanup({"task": {"id": 28}}) + + assert not path.exists(folder) + assert not path.exists(f"{folder}.tar.gz") + + @pytest.mark.asyncio async def test_cleanup_async(async_code_resource): makedirs(folder, exist_ok=True) @@ -122,7 +144,7 @@ def test_download_invalid_token(code_resource): "/code/download", json={"task": {"id": 28, "token": "invalid"}}, headers={ - "Accept": "application/gzip", + "Accept": "application/json", "Content-Type": "application/json", }, ) @@ -159,7 +181,41 @@ async def test_download_async_invalid_token(async_code_resource): "/code/download", json={"task": {"id": 28, "token": "invalid"}}, headers={ - "Accept": "application/gzip", + "Accept": "application/json", + "Content-Type": "application/json", + }, + ) + + # Does not download code + assert not path.exists(folder) + + +def test_download_missing_content_type(code_resource): + # Mock client response + response_mock = MagicMock() + response_mock.status_code = 200 + response_mock.is_error = False + response_mock.headers = {} + response_mock.iter_bytes.return_value = iter([b""]) + + code_resource._client._client.stream.return_value.__enter__.return_value = ( + response_mock + ) + + # Call download + with pytest.raises( + Exception, + match="Unexpected Content-Type: while downloading code.", + ): + code_resource.download({"task": {"id": 28, "token": "abcdef"}}) + + # Hits the API + code_resource._client._client.stream.assert_called_once_with( + "post", + "/code/download", + json={"task": {"id": 28, "token": "abcdef"}}, + headers={ + "Accept": "application/json", "Content-Type": "application/json", }, ) @@ -168,12 +224,53 @@ async def test_download_async_invalid_token(async_code_resource): assert not path.exists(folder) -def test_download(fixture_tarfile, code_resource): +@pytest.mark.asyncio +async def test_download_missing_content_type_async(async_code_resource): # Mock client response response_mock = MagicMock() response_mock.status_code = 200 response_mock.is_error = False - response_mock.headers = {"x-automa-proposal-token": "ghijkl"} + response_mock.headers = {} + + response_byte_reader = AsyncMock() + response_byte_reader.__aiter__.return_value = iter([b""]) + response_mock.aiter_bytes.return_value = response_byte_reader + + async_code_resource._client._client.stream.return_value.__aenter__.return_value = ( + response_mock + ) + + # Call download + with pytest.raises( + Exception, + match="Unexpected Content-Type: while downloading code.", + ): + await async_code_resource.download({"task": {"id": 28, "token": "abcdef"}}) + + # Hits the API + async_code_resource._client._client.stream.assert_called_once_with( + "post", + "/code/download", + json={"task": {"id": 28, "token": "abcdef"}}, + headers={ + "Accept": "application/json", + "Content-Type": "application/json", + }, + ) + + # Does not download code + assert not path.exists(folder) + + +def test_download_proxy(fixture_tarfile, code_resource): + # Mock client response + response_mock = MagicMock() + response_mock.status_code = 200 + response_mock.is_error = False + response_mock.headers = { + "x-automa-proposal-token": "ghijkl", + "Content-Type": "application/gzip", + } with open(fixture_tarfile, "rb") as f: response_mock.iter_bytes.return_value = iter([f.read()]) @@ -193,7 +290,7 @@ def test_download(fixture_tarfile, code_resource): "/code/download", json={"task": {"id": 28, "token": "abcdef"}}, headers={ - "Accept": "application/gzip", + "Accept": "application/json", "Content-Type": "application/json", }, ) @@ -213,12 +310,15 @@ def test_download(fixture_tarfile, code_resource): @pytest.mark.asyncio -async def test_download_async(fixture_tarfile, async_code_resource): +async def test_download_proxy_async(fixture_tarfile, async_code_resource): # Mock client response response_mock = MagicMock() response_mock.status_code = 200 response_mock.is_error = False - response_mock.headers = {"x-automa-proposal-token": "ghijkl"} + response_mock.headers = { + "x-automa-proposal-token": "ghijkl", + "Content-Type": "application/gzip", + } with open(fixture_tarfile, "rb") as f: response_byte_reader = AsyncMock() @@ -243,7 +343,7 @@ async def test_download_async(fixture_tarfile, async_code_resource): "/code/download", json={"task": {"id": 28, "token": "abcdef"}}, headers={ - "Accept": "application/gzip", + "Accept": "application/json", "Content-Type": "application/json", }, ) @@ -261,7 +361,7 @@ async def test_download_async(fixture_tarfile, async_code_resource): def test_propose_no_token(fixture_tarfile, code_resource): - test_download(fixture_tarfile, code_resource) + test_download_proxy(fixture_tarfile, code_resource) remove(proposal_token_file) @@ -274,7 +374,7 @@ def test_propose_no_token(fixture_tarfile, code_resource): @pytest.mark.asyncio async def test_propose_async_no_token(fixture_tarfile, async_code_resource): - await test_download_async(fixture_tarfile, async_code_resource) + await test_download_proxy_async(fixture_tarfile, async_code_resource) remove(proposal_token_file) @@ -286,7 +386,7 @@ async def test_propose_async_no_token(fixture_tarfile, async_code_resource): def test_propose_invalid_token(fixture_tarfile, code_resource): - test_download(fixture_tarfile, code_resource) + test_download_proxy(fixture_tarfile, code_resource) with open(f"{folder}/README.md", "w") as f: f.write("Content\n") @@ -322,7 +422,7 @@ def test_propose_invalid_token(fixture_tarfile, code_resource): @pytest.mark.asyncio async def test_propose_async_invalid_token(fixture_tarfile, async_code_resource): - await test_download_async(fixture_tarfile, async_code_resource) + await test_download_proxy_async(fixture_tarfile, async_code_resource) with open(f"{folder}/README.md", "w") as f: f.write("Content\n") @@ -356,8 +456,8 @@ async def test_propose_async_invalid_token(fixture_tarfile, async_code_resource) ) -def test_propose(fixture_tarfile, code_resource): - test_download(fixture_tarfile, code_resource) +def test_propose_proxy(fixture_tarfile, code_resource): + test_download_proxy(fixture_tarfile, code_resource) with open(f"{folder}/README.md", "w") as f: f.write("Content\n") @@ -390,8 +490,8 @@ def test_propose(fixture_tarfile, code_resource): @pytest.mark.asyncio -async def test_propose_async(fixture_tarfile, async_code_resource): - await test_download_async(fixture_tarfile, async_code_resource) +async def test_propose_proxy_async(fixture_tarfile, async_code_resource): + await test_download_proxy_async(fixture_tarfile, async_code_resource) with open(f"{folder}/README.md", "w") as f: f.write("Content\n") @@ -424,7 +524,7 @@ async def test_propose_async(fixture_tarfile, async_code_resource): def test_propose_with_added_files(fixture_tarfile, code_resource): - code_folder = test_download(fixture_tarfile, code_resource) + code_folder = test_download_proxy(fixture_tarfile, code_resource) with open(f"{code_folder.path}/NEW.md", "w") as f: f.write("Content\n") @@ -459,7 +559,7 @@ def test_propose_with_added_files(fixture_tarfile, code_resource): def test_propose_with_added_files_using_add_all(fixture_tarfile, code_resource): - code_folder = test_download(fixture_tarfile, code_resource) + code_folder = test_download_proxy(fixture_tarfile, code_resource) with open(f"{code_folder.path}/NEW.md", "w") as f: f.write("Content\n") @@ -494,7 +594,7 @@ def test_propose_with_added_files_using_add_all(fixture_tarfile, code_resource): def test_propose_with_proposal_properties(fixture_tarfile, code_resource): - test_download(fixture_tarfile, code_resource) + test_download_proxy(fixture_tarfile, code_resource) with open(f"{folder}/README.md", "w") as f: f.write("Content\n") @@ -534,7 +634,7 @@ def test_propose_with_proposal_properties(fixture_tarfile, code_resource): def test_propose_with_metadata(fixture_tarfile, code_resource): - test_download(fixture_tarfile, code_resource) + test_download_proxy(fixture_tarfile, code_resource) with open(f"{folder}/README.md", "w") as f: f.write("Content\n") @@ -570,3 +670,178 @@ def test_propose_with_metadata(fixture_tarfile, code_resource): "Content-Type": "application/json", }, ) + + +def test_download_direct(fixture_git, code_resource): + # Mock client response + response_mock = MagicMock() + response_mock.status_code = 200 + response_mock.is_error = False + response_mock.headers = { + "x-automa-proposal-token": "ghijkl", + "Content-Type": "application/json", + } + response_mock.iter_bytes.return_value = iter( + [f'{{"type":"direct","url":"file://{fixture_git}"}}'.encode()] + ) + + code_resource._client._client.stream.return_value.__enter__.return_value = ( + response_mock + ) + + # Call download + created_folder = code_resource.download({"task": {"id": 28, "token": "abcdef"}}) + + # Returns path to downloaded code + assert created_folder.path == folder + + # Hits the API + code_resource._client._client.stream.assert_called_once_with( + "post", + "/code/download", + json={"task": {"id": 28, "token": "abcdef"}}, + headers={ + "Accept": "application/json", + "Content-Type": "application/json", + }, + ) + + # Clones the repo + assert path.exists(folder) + assert sorted(listdir(folder)) == [ + ".git", + "README.md", + ] + + # Saves proposal token + with open(proposal_token_file, "r") as f: + assert f.read() == "ghijkl" + + # Saves the base commit + with open(proposal_base_commit_file, "r") as f: + assert f.read() == "cc3f46ae7fdf71747b66b3e4272c0e5fe290d116" + + +@pytest.mark.asyncio +async def test_download_direct_async(fixture_git, async_code_resource): + # Mock client response + response_mock = MagicMock() + response_mock.status_code = 200 + response_mock.is_error = False + response_mock.headers = { + "x-automa-proposal-token": "ghijkl", + "Content-Type": "application/json", + } + + response_byte_reader = AsyncMock() + response_byte_reader.__aiter__.return_value = iter( + [f'{{"type":"direct","url":"file://{fixture_git}"}}'.encode()] + ) + response_mock.aiter_bytes.return_value = response_byte_reader + + async_code_resource._client._client.stream.return_value.__aenter__.return_value = ( + response_mock + ) + + # Call download + created_folder = await async_code_resource.download( + {"task": {"id": 28, "token": "abcdef"}} + ) + + # Returns path to downloaded code + assert created_folder.path == folder + + # Hits the API + async_code_resource._client._client.stream.assert_called_once_with( + "post", + "/code/download", + json={"task": {"id": 28, "token": "abcdef"}}, + headers={ + "Accept": "application/json", + "Content-Type": "application/json", + }, + ) + + # Clones the repo + assert path.exists(folder) + assert sorted(listdir(folder)) == [ + ".git", + "README.md", + ] + + # Saves proposal token + with open(proposal_token_file, "r") as f: + assert f.read() == "ghijkl" + + # Saves the base commit + with open(proposal_base_commit_file, "r") as f: + assert f.read() == "cc3f46ae7fdf71747b66b3e4272c0e5fe290d116" + + +def test_propose_direct(fixture_git, code_resource): + test_download_direct(fixture_git, code_resource) + + with open(f"{folder}/README.md", "w") as f: + f.write("Content\n") + + # Mock client response + response_mock = MagicMock() + response_mock.status_code = 204 + response_mock.is_error = False + + code_resource._client._client.request.return_value = response_mock + + code_resource.propose({"task": {"id": 28, "token": "abcdef"}}) + + # Hits the API + code_resource._client._client.request.assert_called_once_with( + "post", + "/code/propose", + json={ + "task": {"id": 28, "token": "abcdef"}, + "proposal": { + "token": "ghijkl", + "diff": "diff --git a/README.md b/README.md\nindex e69de29..39c9f36 100644\n--- a/README.md\n+++ b/README.md\n@@ -0,0 +1 @@\n+Content\n", + "base_commit": "cc3f46ae7fdf71747b66b3e4272c0e5fe290d116", + }, + }, + headers={ + "Accept": "application/json", + "Content-Type": "application/json", + }, + ) + + +@pytest.mark.asyncio +async def test_propose_direct_async(fixture_git, async_code_resource): + await test_download_direct_async(fixture_git, async_code_resource) + + with open(f"{folder}/README.md", "w") as f: + f.write("Content\n") + + # Mock client response + response_mock = MagicMock() + response_mock.status_code = 204 + response_mock.is_error = False + + async_code_resource._client._client.request = AsyncMock(return_value=response_mock) + + await async_code_resource.propose({"task": {"id": 28, "token": "abcdef"}}) + + # Hits the API + async_code_resource._client._client.request.assert_called_once_with( + "post", + "/code/propose", + json={ + "task": {"id": 28, "token": "abcdef"}, + "proposal": { + "token": "ghijkl", + "diff": "diff --git a/README.md b/README.md\nindex e69de29..39c9f36 100644\n--- a/README.md\n+++ b/README.md\n@@ -0,0 +1 @@\n+Content\n", + "base_commit": "cc3f46ae7fdf71747b66b3e4272c0e5fe290d116", + }, + }, + headers={ + "Accept": "application/json", + "Content-Type": "application/json", + }, + ) diff --git a/uv.lock b/uv.lock index 04bbaf2..caf831a 100644 --- a/uv.lock +++ b/uv.lock @@ -31,7 +31,7 @@ wheels = [ [[package]] name = "automa-bot" -version = "0.2.3" +version = "0.3.0" source = { virtual = "packages/bot" } dependencies = [ { name = "anyio" },