wehub-resource-sync
79 lines
2.3 KiB
Python
79 lines
2.3 KiB
Python
from unittest import mock
|
|
|
|
import git
|
|
import pytest
|
|
|
|
from mlflow.tracking.context.git_context import GitRunContext
|
|
from mlflow.utils.mlflow_tags import (
|
|
MLFLOW_GIT_BRANCH,
|
|
MLFLOW_GIT_COMMIT,
|
|
MLFLOW_GIT_REPO_URL,
|
|
)
|
|
|
|
MOCK_SCRIPT_NAME = "/path/to/script.py"
|
|
MOCK_COMMIT_HASH = "commit-hash"
|
|
MOCK_BRANCH_NAME = "feature/test"
|
|
MOCK_REPO_URL = "git@github.com:mlflow/mlflow.git"
|
|
|
|
|
|
@pytest.fixture
|
|
def patch_script_name():
|
|
patch_sys_argv = mock.patch("sys.argv", [MOCK_SCRIPT_NAME])
|
|
patch_os_path_isfile = mock.patch("os.path.isfile", return_value=False)
|
|
with patch_sys_argv, patch_os_path_isfile:
|
|
yield
|
|
|
|
|
|
@pytest.fixture
|
|
def patch_git_repo():
|
|
mock_repo = mock.Mock()
|
|
mock_repo.head.commit.hexsha = MOCK_COMMIT_HASH
|
|
mock_repo.active_branch.name = MOCK_BRANCH_NAME
|
|
mock_repo.ignored.return_value = []
|
|
mock_remote = mock.Mock()
|
|
mock_remote.url = MOCK_REPO_URL
|
|
mock_repo.remotes = [mock_remote]
|
|
with mock.patch("git.Repo", return_value=mock_repo):
|
|
yield mock_repo
|
|
|
|
|
|
def test_git_run_context_in_context_true(patch_script_name, patch_git_repo):
|
|
assert GitRunContext().in_context()
|
|
|
|
|
|
def test_git_run_context_in_context_false(patch_script_name):
|
|
with mock.patch("git.Repo", side_effect=git.InvalidGitRepositoryError):
|
|
assert not GitRunContext().in_context()
|
|
|
|
|
|
def test_git_run_context_tags(patch_script_name, patch_git_repo):
|
|
assert GitRunContext().tags() == {
|
|
MLFLOW_GIT_COMMIT: MOCK_COMMIT_HASH,
|
|
MLFLOW_GIT_BRANCH: MOCK_BRANCH_NAME,
|
|
MLFLOW_GIT_REPO_URL: MOCK_REPO_URL,
|
|
}
|
|
|
|
|
|
def test_git_run_context_tags_omits_missing_fields(patch_script_name):
|
|
mock_repo = mock.Mock()
|
|
mock_repo.head.commit.hexsha = MOCK_COMMIT_HASH
|
|
mock_repo.active_branch.name = MOCK_BRANCH_NAME
|
|
mock_repo.ignored.return_value = []
|
|
mock_repo.remotes = []
|
|
with mock.patch("git.Repo", return_value=mock_repo):
|
|
assert GitRunContext().tags() == {
|
|
MLFLOW_GIT_COMMIT: MOCK_COMMIT_HASH,
|
|
MLFLOW_GIT_BRANCH: MOCK_BRANCH_NAME,
|
|
}
|
|
|
|
|
|
def test_git_run_context_caching(patch_script_name):
|
|
with mock.patch("git.Repo") as mock_repo:
|
|
context = GitRunContext()
|
|
context.in_context()
|
|
context.tags()
|
|
assert mock_repo.call_count == 3
|
|
context.in_context()
|
|
context.tags()
|
|
assert mock_repo.call_count == 3
|