build_tools/benchmarks/post_benchmark_comment_test.py

#!/usr/bin/env python3
# Copyright 2022 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

import http.client
import requests
import unittest
from unittest import mock
from typing import Any

import post_benchmark_comment


class GithubClientTest(unittest.TestCase):
    def setUp(self):
        self._mock_response = mock.create_autospec(requests.Response)
        self._mock_requester = mock.create_autospec(post_benchmark_comment.APIRequester)
        self._mock_requester.get.return_value = self._mock_response
        self._mock_requester.post.return_value = self._mock_response
        self._mock_requester.patch.return_value = self._mock_response

    def test_post_to_gist(self):
        gist_url = "https://example.com/123455/1234.md"
        self._mock_response.status_code = http.client.CREATED
        self._mock_response.json.return_value = {
            "html_url": gist_url,
            "truncated": False,
        }
        client = post_benchmark_comment.GithubClient(self._mock_requester)

        url = client.post_to_gist(filename="1234.md", content="xyz")

        self.assertEqual(url, gist_url)
        self._mock_requester.post.assert_called_once_with(
            endpoint=post_benchmark_comment.GITHUB_GIST_API,
            payload={"public": True, "files": {"1234.md": {"content": "xyz"}}},
        )

    def test_post_to_gist_truncated(self):
        gist_url = "example.com/123455/1234.md"
        self._mock_response.status_code = http.client.CREATED
        self._mock_response.json.return_value = {
            "html_url": gist_url,
            "truncated": True,
        }
        client = post_benchmark_comment.GithubClient(self._mock_requester)

        with self.assertRaises(RuntimeError) as _:
            client.post_to_gist(filename="1234.md", content="xyz")

    def test_get_previous_comment_on_pr(self):
        first_mock_response = mock.create_autospec(requests.Response)
        first_mock_response.status_code = http.client.OK
        first_mock_response.json.return_value = [
            {"id": 1, "user": {"login": "bot"}, "body": "comment id: abcd"},
            {"id": 2, "user": {"login": "user"}, "body": "comment id: 1234"},
        ]
        second_mock_response = mock.create_autospec(requests.Response)
        second_mock_response.status_code = http.client.OK
        second_mock_response.json.return_value = [
            {"id": 3, "user": {"login": "bot"}, "body": "comment id: 1234"}
        ]
        mock_requester = mock.create_autospec(post_benchmark_comment.APIRequester)
        mock_requester.get.side_effect = [first_mock_response, second_mock_response]
        client = post_benchmark_comment.GithubClient(mock_requester)

        comment_id = client.get_previous_comment_on_pr(
            pr_number=23,
            comment_bot_user="bot",
            comment_type_id="1234",
            query_comment_per_page=2,
            max_pages_to_search=10,
        )

        self.assertEqual(comment_id, 3)
        self.assertEqual(mock_requester.get.call_count, 2)
        endpoint_url = (
            f"{post_benchmark_comment.GITHUB_IREE_API_PREFIX}/issues/23/comments"
        )
        mock_requester.get.assert_any_call(
            endpoint=endpoint_url,
            payload={"per_page": 2, "page": 1, "sort": "updated", "direction": "desc"},
        )
        mock_requester.get.assert_any_call(
            endpoint=endpoint_url,
            payload={"per_page": 2, "page": 2, "sort": "updated", "direction": "desc"},
        )

    def test_get_previous_comment_on_pr_not_found(self):
        mock_response = mock.create_autospec(requests.Response)
        mock_response.status_code = http.client.OK
        mock_response.json.return_value = [
            {"id": 1, "user": {"login": "bot"}, "body": "comment id: 5678"}
        ]
        mock_requester = mock.create_autospec(post_benchmark_comment.APIRequester)
        mock_requester.get.side_effect = [mock_response] * 10
        client = post_benchmark_comment.GithubClient(mock_requester)

        comment_id = client.get_previous_comment_on_pr(
            pr_number=23,
            comment_bot_user="bot",
            comment_type_id="1234",
            query_comment_per_page=1,
            max_pages_to_search=10,
        )

        self.assertIsNone(comment_id)
        self.assertEqual(mock_requester.get.call_count, 10)
        endpoint_url = (
            f"{post_benchmark_comment.GITHUB_IREE_API_PREFIX}/issues/23/comments"
        )
        mock_requester.get.assert_any_call(
            endpoint=endpoint_url,
            payload={"per_page": 1, "page": 1, "sort": "updated", "direction": "desc"},
        )
        mock_requester.get.assert_any_call(
            endpoint=endpoint_url,
            payload={"per_page": 1, "page": 10, "sort": "updated", "direction": "desc"},
        )

    def test_update_comment_on_pr(self):
        self._mock_response.status_code = http.client.OK
        client = post_benchmark_comment.GithubClient(self._mock_requester)

        client.update_comment_on_pr(comment_id=123, content="xyz")

        self._mock_requester.patch.assert_called_once_with(
            endpoint=f"{post_benchmark_comment.GITHUB_IREE_API_PREFIX}/issues/comments/123",
            payload={"body": "xyz"},
        )

    def test_create_comment_on_pr(self):
        self._mock_response.status_code = http.client.CREATED
        client = post_benchmark_comment.GithubClient(self._mock_requester)

        client.create_comment_on_pr(pr_number=1234, content="xyz")

        self._mock_requester.post.assert_called_once_with(
            endpoint=f"{post_benchmark_comment.GITHUB_IREE_API_PREFIX}/issues/1234/comments",
            payload={"body": "xyz"},
        )

    def test_get_pull_request_head_commit(self):
        self._mock_response.status_code = http.client.OK
        self._mock_response.json.return_value = {"head": {"sha": "sha123"}}
        client = post_benchmark_comment.GithubClient(self._mock_requester)

        commit_sha = client.get_pull_request_head_commit(pr_number=123)

        self.assertEqual(commit_sha, "sha123")
        self._mock_requester.get.assert_called_once_with(
            endpoint=f"{post_benchmark_comment.GITHUB_IREE_API_PREFIX}/pulls/123"
        )


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