diff --git a/.generator/cli.py b/.generator/cli.py index 0fd97d11c6fe..adc3cc4c242c 100644 --- a/.generator/cli.py +++ b/.generator/cli.py @@ -24,6 +24,7 @@ LIBRARIAN_DIR = "librarian" GENERATE_REQUEST_FILE = "generate-request.json" SOURCE_DIR = "source" +OUTPUT_DIR = "output" def _read_json_file(path): @@ -97,7 +98,6 @@ def _build_bazel_target(bazel_rule): subprocess.run( command, cwd=f"{SOURCE_DIR}/googleapis", - capture_output=True, text=True, check=True, ) @@ -106,6 +106,53 @@ def _build_bazel_target(bazel_rule): raise ValueError(f"Bazel build for {bazel_rule} rule failed.") from e +def _locate_and_extract_artifact(bazel_rule: str, library_id: str): + """Finds and extracts the tarball artifact from a Bazel build. + + Args: + bazel_rule (str): The Bazel rule that was built. + library_id (str): The ID of the library being generated. + + Raises: + ValueError: If failed to locate or extract artifact. + """ + try: + # 1. Find the bazel-bin output directory. + logger.info("Locating Bazel output directory...") + info_command = ["bazelisk", "info", "bazel-bin"] + result = subprocess.run( + info_command, + cwd=f"{SOURCE_DIR}/googleapis", + text=True, + check=True, + capture_output=True, + ) + bazel_bin_path = result.stdout.strip() + + # 2. Construct the path to the generated tarball. + rule_path, rule_name = bazel_rule.split(":") + tarball_name = f"{rule_name}.tar.gz" + tarball_path = os.path.join(bazel_bin_path, rule_path.strip("/"), tarball_name) + logger.info(f"Found artifact at: {tarball_path}") + + # 3. Create a staging directory. + staging_dir = os.path.join(OUTPUT_DIR, "owl-bot-staging", library_id) + os.makedirs(staging_dir, exist_ok=True) + logger.info(f"Preparing staging directory: {staging_dir}") + + # 4. Extract the artifact. + extract_command = ["tar", "-xvf", tarball_path, "--strip-components=1"] + subprocess.run( + extract_command, cwd=staging_dir, capture_output=True, text=True, check=True + ) + logger.info(f"Artifact {tarball_path} extracted successfully.") + + except Exception as e: + raise ValueError( + f"Failed to locate or extract artifact for {bazel_rule} rule" + ) from e + + def handle_generate(): """The main coordinator for the code generation process. @@ -129,6 +176,7 @@ def handle_generate(): if api_path: bazel_rule = _determine_bazel_rule(api_path) _build_bazel_target(bazel_rule) + _locate_and_extract_artifact(bazel_rule, library_id) logger.info(json.dumps(request_data, indent=2)) except Exception as e: diff --git a/.generator/test_cli.py b/.generator/test_cli.py index a5cd14604c04..f2c4ebe4e706 100644 --- a/.generator/test_cli.py +++ b/.generator/test_cli.py @@ -24,6 +24,7 @@ _read_json_file, _determine_bazel_rule, _build_bazel_target, + _locate_and_extract_artifact, handle_generate, handle_build, handle_configure, @@ -118,6 +119,52 @@ def test_determine_bazel_rule_command_fails(mocker, caplog): assert "Found Bazel rule" not in caplog.text +def test_locate_and_extract_artifact_success(mocker, caplog): + """ + Tests that the artifact helper calls the correct sequence of commands. + """ + caplog.set_level(logging.INFO) + mock_info_result = MagicMock(stdout="/path/to/bazel-bin\n") + mock_tar_result = MagicMock(returncode=0) + mocker.patch("cli.subprocess.run", side_effect=[mock_info_result, mock_tar_result]) + mock_makedirs = mocker.patch("cli.os.makedirs") + _locate_and_extract_artifact( + "//google/cloud/language/v1:rule-py", + "google-cloud-language", + ) + + assert ( + "Found artifact at: /path/to/bazel-bin/google/cloud/language/v1/rule-py.tar.gz" + in caplog.text + ) + assert ( + "Preparing staging directory: output/owl-bot-staging/google-cloud-language" + in caplog.text + ) + assert ( + "Artifact /path/to/bazel-bin/google/cloud/language/v1/rule-py.tar.gz extracted successfully" + in caplog.text + ) + mock_makedirs.assert_called_once() + + +def test_locate_and_extract_artifact_fails(mocker, caplog): + """ + Tests that an exception is raised if the subprocess command fails. + """ + caplog.set_level(logging.INFO) + mocker.patch( + "cli.subprocess.run", + side_effect=subprocess.CalledProcessError(1, "cmd", stderr="Bazel error"), + ) + + with pytest.raises(ValueError): + _locate_and_extract_artifact( + "//google/cloud/language/v1:rule-py", + "google-cloud-language", + ) + + def test_handle_generate_success(caplog, mock_generate_request_file, mocker): """ Tests the successful execution path of handle_generate. @@ -129,6 +176,8 @@ def test_handle_generate_success(caplog, mock_generate_request_file, mocker): ) mock_build_target = mocker.patch("cli._build_bazel_target") + mock_locate_and_extract_artifact = mocker.patch("cli._locate_and_extract_artifact") + handle_generate() mock_determine_rule.assert_called_once_with("google/cloud/language/v1")