pytorch

1
"""test_check_labels.py"""
2

3
from typing import Any, List
4
from unittest import main, mock, TestCase
5

6
from check_labels import (
7
    add_label_err_comment,
8
    delete_all_label_err_comments,
9
    main as check_labels_main,
10
)
11
from github_utils import GitHubComment
12
from label_utils import BOT_AUTHORS, LABEL_ERR_MSG_TITLE
13
from test_trymerge import mock_gh_get_info, mocked_gh_graphql
14
from trymerge import GitHubPR
15

16

17
def mock_parse_args() -> object:
18
    class Object:
19
        def __init__(self) -> None:
20
            self.pr_num = 76123
21
            self.exit_non_zero = False
22

23
    return Object()
24

25

26
def mock_add_label_err_comment(pr: "GitHubPR") -> None:
27
    pass
28

29

30
def mock_delete_all_label_err_comments(pr: "GitHubPR") -> None:
31
    pass
32

33

34
def mock_get_comments() -> List[GitHubComment]:
35
    return [
36
        # Case 1 - a non label err comment
37
        GitHubComment(
38
            body_text="mock_body_text",
39
            created_at="",
40
            author_login="",
41
            author_association="",
42
            editor_login=None,
43
            database_id=1,
44
            url="",
45
        ),
46
        # Case 2 - a label err comment
47
        GitHubComment(
48
            body_text=" #" + LABEL_ERR_MSG_TITLE.replace("`", ""),
49
            created_at="",
50
            author_login=BOT_AUTHORS[1],
51
            author_association="",
52
            editor_login=None,
53
            database_id=2,
54
            url="",
55
        ),
56
    ]
57

58

59
class TestCheckLabels(TestCase):
60
    @mock.patch("trymerge.gh_graphql", side_effect=mocked_gh_graphql)
61
    @mock.patch("trymerge.GitHubPR.get_comments", return_value=[mock_get_comments()[0]])
62
    @mock.patch("check_labels.gh_post_pr_comment")
63
    def test_correctly_add_label_err_comment(
64
        self, mock_gh_post_pr_comment: Any, mock_get_comments: Any, mock_gh_grphql: Any
65
    ) -> None:
66
        "Test add label err comment when similar comments don't exist."
67
        pr = GitHubPR("pytorch", "pytorch", 75095)
68
        add_label_err_comment(pr)
69
        mock_gh_post_pr_comment.assert_called_once()
70

71
    @mock.patch("trymerge.gh_graphql", side_effect=mocked_gh_graphql)
72
    @mock.patch("trymerge.GitHubPR.get_comments", return_value=[mock_get_comments()[1]])
73
    @mock.patch("check_labels.gh_post_pr_comment")
74
    def test_not_add_label_err_comment(
75
        self, mock_gh_post_pr_comment: Any, mock_get_comments: Any, mock_gh_grphql: Any
76
    ) -> None:
77
        "Test not add label err comment when similar comments exist."
78
        pr = GitHubPR("pytorch", "pytorch", 75095)
79
        add_label_err_comment(pr)
80
        mock_gh_post_pr_comment.assert_not_called()
81

82
    @mock.patch("trymerge.gh_graphql", side_effect=mocked_gh_graphql)
83
    @mock.patch("trymerge.GitHubPR.get_comments", return_value=mock_get_comments())
84
    @mock.patch("check_labels.gh_delete_comment")
85
    def test_correctly_delete_all_label_err_comments(
86
        self, mock_gh_delete_comment: Any, mock_get_comments: Any, mock_gh_grphql: Any
87
    ) -> None:
88
        "Test only delete label err comment."
89
        pr = GitHubPR("pytorch", "pytorch", 75095)
90
        delete_all_label_err_comments(pr)
91
        mock_gh_delete_comment.assert_called_once_with("pytorch", "pytorch", 2)
92

93
    @mock.patch("trymerge.gh_get_pr_info", return_value=mock_gh_get_info())
94
    @mock.patch("check_labels.parse_args", return_value=mock_parse_args())
95
    @mock.patch("check_labels.has_required_labels", return_value=False)
96
    @mock.patch(
97
        "check_labels.delete_all_label_err_comments",
98
        side_effect=mock_delete_all_label_err_comments,
99
    )
100
    @mock.patch(
101
        "check_labels.add_label_err_comment", side_effect=mock_add_label_err_comment
102
    )
103
    def test_ci_comments_and_exit0_without_required_labels(
104
        self,
105
        mock_add_label_err_comment: Any,
106
        mock_delete_all_label_err_comments: Any,
107
        mock_has_required_labels: Any,
108
        mock_parse_args: Any,
109
        mock_gh_get_info: Any,
110
    ) -> None:
111
        with self.assertRaises(SystemExit) as sys_exit:
112
            check_labels_main()
113
        self.assertEqual(str(sys_exit.exception), "0")
114
        mock_add_label_err_comment.assert_called_once()
115
        mock_delete_all_label_err_comments.assert_not_called()
116

117
    @mock.patch("trymerge.gh_get_pr_info", return_value=mock_gh_get_info())
118
    @mock.patch("check_labels.parse_args", return_value=mock_parse_args())
119
    @mock.patch("check_labels.has_required_labels", return_value=True)
120
    @mock.patch(
121
        "check_labels.delete_all_label_err_comments",
122
        side_effect=mock_delete_all_label_err_comments,
123
    )
124
    @mock.patch(
125
        "check_labels.add_label_err_comment", side_effect=mock_add_label_err_comment
126
    )
127
    def test_ci_exit0_with_required_labels(
128
        self,
129
        mock_add_label_err_comment: Any,
130
        mock_delete_all_label_err_comments: Any,
131
        mock_has_required_labels: Any,
132
        mock_parse_args: Any,
133
        mock_gh_get_info: Any,
134
    ) -> None:
135
        with self.assertRaises(SystemExit) as sys_exit:
136
            check_labels_main()
137
        self.assertEqual(str(sys_exit.exception), "0")
138
        mock_add_label_err_comment.assert_not_called()
139
        mock_delete_all_label_err_comments.assert_called_once()
140

141

142
if __name__ == "__main__":
143
    main()
144