# Copyright 2025 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Tests for buildbucket tools."""

import json
import os
import pathlib
import subprocess
import sys
import unittest
from unittest import mock

sys.path.insert(
    0,
    os.path.abspath(
        pathlib.Path(__file__)
        .resolve()
        .parent.parent.joinpath(pathlib.Path("infra_lib"))
    ),
)
import buildbucket


class BuildbucketTest(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        self.mock_context = mock.AsyncMock()
        self.mock_context.info = mock.AsyncMock()

    @mock.patch("subprocess.run")
    async def test_get_build_status_success(self, mock_subprocess_run):
        build_id = "12345"
        expected_status = "SUCCESS"
        mock_subprocess_run.return_value = subprocess.CompletedProcess(
            args=[],
            returncode=0,
            stdout=json.dumps({"status": expected_status}),
            stderr="",
        )

        status = await buildbucket.get_build_status(
            self.mock_context,
            build_id,
        )

        self.assertEqual(status, expected_status)
        expected_command = [
            "prpc",
            "call",
            "cr-buildbucket.appspot.com",
            "buildbucket.v2.Builds.GetBuildStatus",
        ]
        mock_subprocess_run.assert_called_once_with(
            expected_command,
            capture_output=True,
            input=json.dumps({"id": build_id}),
            check=True,
            text=True,
        )

    @mock.patch("subprocess.run")
    async def test_get_build_status_exception(self, mock_subprocess_run):
        build_id = "12345"
        mock_subprocess_run.side_effect = Exception("PRPC call failed")

        with self.assertRaisesRegex(Exception, "PRPC call failed"):
            await buildbucket.get_build_status(self.mock_context, build_id)

    @mock.patch("subprocess.run")
    async def test_get_build_from_id_success(self, mock_subprocess_run):
        build_id = "12345"
        fields = ["steps", "tags"]
        expected_output = '{"id": "12345", "steps": [], "tags": []}'
        mock_subprocess_run.return_value = subprocess.CompletedProcess(
            args=[], returncode=0, stdout=expected_output, stderr=""
        )

        output = await buildbucket.get_build_from_id(
            self.mock_context,
            build_id,
            fields,
        )

        self.assertEqual(output, expected_output)
        expected_command = [
            "prpc",
            "call",
            "cr-buildbucket.appspot.com",
            "buildbucket.v2.Builds.GetBuild",
        ]
        expected_request = {"id": build_id, "mask": {"fields": "steps,tags"}}
        mock_subprocess_run.assert_called_once_with(
            expected_command,
            capture_output=True,
            input=json.dumps(expected_request),
            check=True,
            text=True,
        )

    @mock.patch("subprocess.run")
    async def test_get_build_from_id_exception(self, mock_subprocess_run):
        build_id = "12345"
        fields = ["steps"]
        mock_subprocess_run.side_effect = Exception("PRPC call failed")

        with self.assertRaisesRegex(Exception, "PRPC call failed"):
            await buildbucket.get_build_from_id(
                self.mock_context,
                build_id,
                fields,
            )

    @mock.patch("subprocess.run")
    async def test_get_build_from_build_number_success(
        self,
        mock_subprocess_run,
    ):
        build_number = 987
        builder_name = "test_builder"
        builder_bucket = "try"
        builder_project = "chromium"
        fields = ["status"]
        expected_output = '{"status": "SUCCESS"}'
        mock_subprocess_run.return_value = subprocess.CompletedProcess(
            args=[], returncode=0, stdout=expected_output, stderr=""
        )

        output = await buildbucket.get_build_from_build_number(
            self.mock_context,
            build_number,
            builder_name,
            builder_bucket,
            builder_project,
            fields,
        )

        self.assertEqual(output, expected_output)
        expected_command = [
            "prpc",
            "call",
            "cr-buildbucket.appspot.com",
            "buildbucket.v2.Builds.GetBuild",
        ]
        expected_request = {
            "buildNumber": build_number,
            "builder": {
                "builder": builder_name,
                "bucket": builder_bucket,
                "project": builder_project,
            },
            "mask": {"fields": "status"},
        }
        mock_subprocess_run.assert_called_once_with(
            expected_command,
            capture_output=True,
            input=json.dumps(expected_request),
            check=True,
            text=True,
        )

    @mock.patch("subprocess.run")
    async def test_get_build_from_build_number_exception(
        self, mock_subprocess_run
    ):
        build_number = 987
        builder_name = "test_builder"
        builder_bucket = "try"
        builder_project = "chromium"
        fields = ["status"]
        mock_subprocess_run.side_effect = Exception("PRPC call failed")

        with self.assertRaisesRegex(Exception, "PRPC call failed"):
            await buildbucket.get_build_from_build_number(
                self.mock_context,
                build_number,
                builder_name,
                builder_bucket,
                builder_project,
                fields,
            )

    @mock.patch("subprocess.run")
    async def test_get_build_success(self, mock_subprocess_run):
        request = {"id": "12345"}
        expected_output = '{"id": "12345"}'
        mock_subprocess_run.return_value = subprocess.CompletedProcess(
            args=[],
            returncode=0,
            stdout=expected_output,
            stderr="",
        )

        output = await buildbucket.get_build(self.mock_context, request)

        self.assertEqual(output, expected_output)
        expected_command = [
            "prpc",
            "call",
            "cr-buildbucket.appspot.com",
            "buildbucket.v2.Builds.GetBuild",
        ]
        mock_subprocess_run.assert_called_once_with(
            expected_command,
            capture_output=True,
            input=json.dumps(request),
            check=True,
            text=True,
        )

    @mock.patch("subprocess.run")
    async def test_get_build_exception(self, mock_subprocess_run):
        request = {"id": "12345"}
        mock_subprocess_run.side_effect = Exception("PRPC call failed")

        with self.assertRaisesRegex(Exception, "PRPC call failed"):
            await buildbucket.get_build(self.mock_context, request)

    @mock.patch("subprocess.run")
    async def test_get_recent_builds_success(self, mock_subprocess_run):
        expected_output = '{"builds": [{"id": "1"}]}'
        mock_subprocess_run.return_value = subprocess.CompletedProcess(
            args=[],
            returncode=0,
            stdout=expected_output,
            stderr="",
        )

        output = await buildbucket.get_recent_builds(
            self.mock_context,
            "test_builder",
            "try",
            "chromium",
            10,
        )

        self.assertEqual(output, expected_output)
        expected_command = [
            "prpc",
            "call",
            "cr-buildbucket.appspot.com",
            "buildbucket.v2.Builds.SearchBuilds",
        ]
        expected_request = {
            "predicate": {
                "builder": {
                    "project": "chromium",
                    "bucket": "try",
                    "builder": "test_builder",
                },
                "status": "ENDED_MASK",
            },
            "page_size": "10",
        }
        mock_subprocess_run.assert_called_once_with(
            expected_command,
            capture_output=True,
            input=json.dumps(expected_request),
            check=True,
            text=True,
        )

    @mock.patch("subprocess.run")
    async def test_get_recent_builds_with_url_encoding_success(
        self, mock_subprocess_run
    ):
        builder_name_encoded = "test%20builder"
        builder_name_decoded = "test builder"
        expected_output = '{"builds": [{"id": "1"}]}'
        mock_subprocess_run.return_value = subprocess.CompletedProcess(
            args=[],
            returncode=0,
            stdout=expected_output,
            stderr="",
        )

        output = await buildbucket.get_recent_builds(
            self.mock_context,
            builder_name_encoded,
            "try",
            "chromium",
            10,
        )

        self.assertEqual(output, expected_output)
        expected_command = [
            "prpc",
            "call",
            "cr-buildbucket.appspot.com",
            "buildbucket.v2.Builds.SearchBuilds",
        ]
        expected_request = {
            "predicate": {
                "builder": {
                    "project": "chromium",
                    "bucket": "try",
                    "builder": builder_name_decoded,
                },
                "status": "ENDED_MASK",
            },
            "page_size": "10",
        }
        mock_subprocess_run.assert_called_once_with(
            expected_command,
            capture_output=True,
            input=json.dumps(expected_request),
            check=True,
            text=True,
        )

    @mock.patch("subprocess.run")
    async def test_get_recent_builds_exception(self, mock_subprocess_run):
        mock_subprocess_run.side_effect = Exception("PRPC call failed")

        with self.assertRaisesRegex(Exception, "PRPC call failed"):
            await buildbucket.get_recent_builds(
                self.mock_context,
                "test_builder",
                "try",
                "chromium",
                10,
            )

    async def test_get_recent_builds_invalid_num_builds(self):
        with self.assertRaisesRegex(
            ValueError, "Provided num_builds 0 is not positive"
        ):
            await buildbucket.get_recent_builds(
                self.mock_context,
                "test_builder",
                "try",
                "chromium",
                0,
            )

    @mock.patch("subprocess.run")
    async def test_get_recent_failed_builds_success(self, mock_subprocess_run):
        expected_output = '{"builds": [{"id": "1", "status": "FAILURE"}]}'
        mock_subprocess_run.return_value = subprocess.CompletedProcess(
            args=[],
            returncode=0,
            stdout=expected_output,
            stderr="",
        )

        output = await buildbucket.get_recent_failed_builds(
            self.mock_context,
            "test_builder",
            "try",
            "chromium",
            10,
        )

        self.assertEqual(output, expected_output)
        expected_command = [
            "prpc",
            "call",
            "cr-buildbucket.appspot.com",
            "buildbucket.v2.Builds.SearchBuilds",
        ]
        expected_request = {
            "predicate": {
                "builder": {
                    "project": "chromium",
                    "bucket": "try",
                    "builder": "test_builder",
                },
                "status": "FAILURE",
            },
            "page_size": "10",
        }
        mock_subprocess_run.assert_called_once_with(
            expected_command,
            capture_output=True,
            input=json.dumps(expected_request),
            check=True,
            text=True,
        )

    @mock.patch("subprocess.run")
    async def test_get_recent_failed_builds_exception(
        self, mock_subprocess_run
    ):
        mock_subprocess_run.side_effect = Exception("PRPC call failed")

        with self.assertRaisesRegex(Exception, "PRPC call failed"):
            await buildbucket.get_recent_failed_builds(
                self.mock_context,
                "test_builder",
                "try",
                "chromium",
                10,
            )

    async def test_get_recent_failed_builds_invalid_num_builds(self):
        with self.assertRaisesRegex(
            ValueError, "Provided num_builds -1 is not positive"
        ):
            await buildbucket.get_recent_failed_builds(
                self.mock_context,
                "test_builder",
                "try",
                "chromium",
                -1,
            )


if __name__ == "__main__":
    unittest.main()
