github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/clients/python-wrapper/tests/utests/test_branch.py

import http

import lakefs_sdk
import pytest

from tests.utests.common import get_test_client, expect_exception_context
from lakefs.repository import Repository
from lakefs.exceptions import ConflictException


def get_test_branch():
    client = get_test_client()
    repo = Repository(repository_id="test_repo", client=client)
    return repo.branch("test_branch")


def test_branch_creation():
    """
    Ensure branches can only be created in repo context
    """
    client = get_test_client()
    repo = Repository(repository_id="test_repo", client=client)
    branch = repo.branch("test_branch")
    assert branch.repo_id == "test_repo"
    assert branch.id == "test_branch"


def test_branch_create(monkeypatch):
    branch = get_test_branch()
    source = "main"
    with monkeypatch.context():
        def monkey_create_branch(_self, repo_name, branch_creation, *_):
            assert repo_name == branch.repo_id
            assert branch_creation.name == branch.id
            assert branch_creation.source == source
            return lakefs_sdk.BranchCreation(name=branch.id, source=source)

        monkeypatch.setattr(lakefs_sdk.BranchesApi, "create_branch", monkey_create_branch)
        branch.create(source)


def test_branch_create_already_exists(monkeypatch):
    branch = get_test_branch()
    source = "main"
    ex = lakefs_sdk.exceptions.ApiException(status=http.HTTPStatus.CONFLICT.value)

    with monkeypatch.context():
        def monkey_create_branch(_self, repo_name, branch_creation, *_):
            raise ex

        monkeypatch.setattr(lakefs_sdk.BranchesApi, "create_branch", monkey_create_branch)

        # Expect success when exist_ok = True
        res = branch.create(source, exist_ok=True)
        assert res.id == branch.id
        assert res.repo_id == branch.repo_id

        # Expect fail on exists
        with expect_exception_context(ConflictException):
            branch.create(source)


def test_branch_head(monkeypatch):
    branch = get_test_branch()
    commit_id = "1234"
    with monkeypatch.context():
        def monkey_get_branch(repo_name, branch_name, *_):
            assert repo_name == branch.repo_id
            assert branch_name == branch.id
            return lakefs_sdk.Ref(commit_id=commit_id, id=branch.id)

        monkeypatch.setattr(branch._client.sdk_client.branches_api, "get_branch", monkey_get_branch)
        res = branch.head
        assert res.id == commit_id
        assert res.repo_id == branch.repo_id


def test_branch_commit(monkeypatch):
    branch = get_test_branch()
    md = {"key": "value"}
    commit_id = "1234"
    commit_message = "test message"
    with monkeypatch.context():
        def monkey_commit(repo_name, branch_name, commits_creation, *_):
            assert repo_name == branch.repo_id
            assert branch_name == branch.id
            assert commits_creation.message == commit_message
            assert commits_creation.metadata == md
            assert commits_creation.allow_empty
            return lakefs_sdk.Commit(
                id=commit_id,
                parents=[""],
                committer="Committer",
                message=commit_message,
                creation_date=123,
                meta_range_id="",
            )

        monkeypatch.setattr(branch._client.sdk_client.commits_api, "commit", monkey_commit)
        res = branch.commit(commit_message, metadata=md, allow_empty=True, ignored_field="test")
        assert res.id == commit_id


def test_branch_delete(monkeypatch):
    branch = get_test_branch()
    with monkeypatch.context():
        def monkey_delete_branch(repo_name, branch_name, *_):
            assert repo_name == branch.repo_id
            assert branch_name == branch.id

        monkeypatch.setattr(branch._client.sdk_client.branches_api, "delete_branch", monkey_delete_branch)
        branch.delete()


def test_branch_revert(monkeypatch):
    branch = get_test_branch()
    ref_id = "ab1234"
    expected_parent = 0
    with monkeypatch.context():
        def monkey_revert_branch(repo_name, branch_name, revert_branch_creation, *_):
            assert repo_name == branch.repo_id
            assert branch_name == branch.id
            assert revert_branch_creation.ref == ref_id
            assert revert_branch_creation.parent_number == expected_parent  # default value

        def monkey_get_commit(repo_name, ref_name, **_):
            assert repo_name == branch.repo_id
            assert ref_name == branch.id
            return lakefs_sdk.Commit(
                id=ref_id,
                parents=[""],
                committer="Committer",
                message="Message",
                creation_date=0,
                meta_range_id=""
            )

        # Test default parent number
        monkeypatch.setattr(branch._client.sdk_client.branches_api, "revert_branch", monkey_revert_branch)
        monkeypatch.setattr(branch._client.sdk_client.commits_api, "get_commit", monkey_get_commit)
        branch.revert(ref_id)
        expected_parent = 2
        # Test set parent number
        branch.revert(ref_id, 2)

        # Test set invalid parent number
        with expect_exception_context(ValueError):
            branch.revert(ref_id, -1)

        expected_parent = 0
        # reference_id passed, but not reference
        with pytest.warns(DeprecationWarning, match="reference_id is deprecated.*"):
            branch.revert(None, reference_id=ref_id)

        # neither reference nor reference_id passed
        with pytest.raises(ValueError, match=".* must be specified"):
            branch.revert(None)

        # both are passed, prefer ``reference_id``
        with pytest.warns(DeprecationWarning, match="reference_id is deprecated.*"):
            # this is not a high-quality test, but it would throw if the revert API
            # was called with reference "hello" due to the monkey-patching above
            # always returning "ab1234" as ref ID.
            c = branch.revert(ref_id, reference_id="hello")
            assert c.id == ref_id