from unittest import TestCase, mock, main
from test_trymerge import mocked_gh_graphql
from trymerge import GitHubPR
from gitutils import get_git_remote_name, get_git_repo_dir, GitRepo
from typing import Any
from tryrebase import rebase_onto, rebase_ghstack_onto

def mocked_rev_parse(branch: str) -> str:
    return branch

class TestRebase(TestCase):

    @mock.patch('trymerge.gh_graphql', side_effect=mocked_gh_graphql)
    @mock.patch('gitutils.GitRepo._run_git')
    @mock.patch('gitutils.GitRepo.rev_parse', side_effect=mocked_rev_parse)
    @mock.patch('tryrebase.gh_post_comment')
    def test_rebase(self, mocked_post_comment: Any, mocked_rp: Any, mocked_run_git: Any, mocked_gql: Any) -> None:
        "Tests rebase successfully"
        pr = GitHubPR("pytorch", "pytorch", 31093)
        repo = GitRepo(get_git_repo_dir(), get_git_remote_name())
        rebase_onto(pr, repo, 'master')
        calls = [mock.call('fetch', 'origin', 'pull/31093/head:pull/31093/head'),
                 mock.call('rebase', 'refs/remotes/origin/master', 'pull/31093/head'),
                 mock.call('push', '-f', 'https://github.com/mingxiaoh/pytorch.git', 'pull/31093/head:master')]
        mocked_run_git.assert_has_calls(calls)
        self.assertTrue(
            "Successfully rebased `master` onto `refs/remotes/origin/master`" in mocked_post_comment.call_args[0][3])

    @mock.patch('trymerge.gh_graphql', side_effect=mocked_gh_graphql)
    @mock.patch('gitutils.GitRepo._run_git')
    @mock.patch('gitutils.GitRepo.rev_parse', side_effect=mocked_rev_parse)
    @mock.patch('tryrebase.gh_post_comment')
    def test_rebase_to_stable(self, mocked_post_comment: Any, mocked_rp: Any, mocked_run_git: Any, mocked_gql: Any) -> None:
        "Tests rebase to viable/strict successfully"
        pr = GitHubPR("pytorch", "pytorch", 31093)
        repo = GitRepo(get_git_repo_dir(), get_git_remote_name())
        rebase_onto(pr, repo, 'viable/strict', False)
        calls = [mock.call('fetch', 'origin', 'pull/31093/head:pull/31093/head'),
                 mock.call('rebase', 'refs/remotes/origin/viable/strict', 'pull/31093/head'),
                 mock.call('push', '-f', 'https://github.com/mingxiaoh/pytorch.git', 'pull/31093/head:master')]
        mocked_run_git.assert_has_calls(calls)
        self.assertTrue(
            "Successfully rebased `master` onto `refs/remotes/origin/viable/strict`" in mocked_post_comment.call_args[0][3])

    @mock.patch('trymerge.gh_graphql', side_effect=mocked_gh_graphql)
    @mock.patch('gitutils.GitRepo._run_git', return_value="Everything up-to-date")
    @mock.patch('gitutils.GitRepo.rev_parse', side_effect=mocked_rev_parse)
    @mock.patch('tryrebase.gh_post_comment')
    def test_no_need_to_rebase(self, mocked_post_comment: Any, mocked_rp: Any, mocked_run_git: Any, mocked_gql: Any) -> None:
        "Tests branch already up to date"
        pr = GitHubPR("pytorch", "pytorch", 31093)
        repo = GitRepo(get_git_repo_dir(), get_git_remote_name())
        rebase_onto(pr, repo, 'master')
        calls = [mock.call('fetch', 'origin', 'pull/31093/head:pull/31093/head'),
                 mock.call('rebase', 'refs/remotes/origin/master', 'pull/31093/head'),
                 mock.call('push', '-f', 'https://github.com/mingxiaoh/pytorch.git', 'pull/31093/head:master')]
        mocked_run_git.assert_has_calls(calls)
        self.assertTrue(
            "Tried to rebase and push PR #31093, but it was already up to date" in mocked_post_comment.call_args[0][3])

    @mock.patch('trymerge.gh_graphql', side_effect=mocked_gh_graphql)
    @mock.patch('gitutils.GitRepo._run_git')
    @mock.patch('gitutils.GitRepo.rev_parse', side_effect=lambda branch: "same sha")
    @mock.patch('tryrebase.gh_post_comment')
    def test_same_sha(self, mocked_post_comment: Any, mocked_rp: Any, mocked_run_git: Any, mocked_gql: Any) -> None:
        "Tests rebase results in same sha"
        pr = GitHubPR("pytorch", "pytorch", 31093)
        repo = GitRepo(get_git_repo_dir(), get_git_remote_name())
        with self.assertRaisesRegex(Exception, 'same sha as the target branch'):
            rebase_onto(pr, repo, 'master')
        with self.assertRaisesRegex(Exception, 'same sha as the target branch'):
            rebase_ghstack_onto(pr, repo, 'master')

if __name__ == "__main__":
    main()