zerto-ai-rewind/tests/test_recover.py

from zerto_rewind_mcp.recover import (
    download_token_from,
    flr_row,
    flr_status,
    session_id_from,
)


def test_session_id_shapes():
    assert session_id_from("abc") == "abc"
    assert session_id_from({"sessionId": "s1"}) == "s1"
    assert session_id_from({"Identifier": "s2"}) == "s2"
    assert session_id_from({"FlrSessionIdentifier": "s3"}) == "s3"


def test_download_token_shapes():
    assert download_token_from("tok") == "tok"
    assert download_token_from({"downloadToken": "t2"}) == "t2"
    assert download_token_from("v1/downloads/abc").startswith("v1/")


def test_flr_list_status():
    payload = [
        {
            "FlrSessionStatus": "MountCompletedSuccessfully",
            "FlrSessionIdentifier": "sess-1",
        }
    ]
    row = flr_row(payload)
    assert flr_status(row).lower() == "mountcompletedsuccessfully"


def test_path_items_shapes():
    from zerto_rewind_mcp.recover import path_items

    assert path_items({"PathItems": [{"Path": "a"}]}) == [{"Path": "a"}]
    assert path_items([{"Path": "b"}]) == [{"Path": "b"}]
    assert path_items(None) == []


def test_decode_flr_path():
    from zerto_rewind_mcp.recover import _decode_flr_path

    assert _decode_flr_path("Volume2-Ext4%2fhome%2fjustin%2fapp-config.yaml") == (
        "Volume2-Ext4/home/justin/app-config.yaml"
    )


def test_resolve_flr_path_picks_browsable_partition():
    import asyncio

    from zerto_rewind_mcp.recover import resolve_flr_path

    root = {
        "PathItems": [
            {"Path": "Volume1-Unknown", "IsBrowsable": False},
            {"Path": "Volume2-Ext4", "IsBrowsable": True},
        ]
    }
    listing = {
        "PathItems": [
            {"Path": "Volume2-Ext4%2fhome%2fjustin%2f.bashrc", "Type": "File"},
            {"Path": "Volume2-Ext4%2fhome%2fjustin%2fapp-config.yaml", "Type": "File"},
        ]
    }

    class FakeClient:
        def __init__(self):
            self.seen = []

        async def browse_flr(self, session_id, path="", recursive=False):
            self.seen.append(path)
            return root if path == "" else listing

    client = FakeClient()
    got = asyncio.run(resolve_flr_path(client, "sess", "/home/justin/app-config.yaml"))
    assert got == "Volume2-Ext4/home/justin/app-config.yaml"
    # must not try the unrestorable partition
    assert "Volume1-Unknown/home/justin" not in client.seen


def test_session_rows_and_id_shapes():
    from zerto_rewind_mcp.recover import session_id_of, session_rows

    assert session_rows([{"a": 1}]) == [{"a": 1}]
    assert session_rows({"a": 1}) == [{"a": 1}]
    assert session_rows(None) == []
    assert session_id_of({"FlrSessionIdentifier": "s1"}) == "s1"
    assert session_id_of({"sessionId": "s2"}) == "s2"
    assert session_id_of({}) == ""


def test_is_live_session():
    from zerto_rewind_mcp.recover import is_live_session

    assert is_live_session({"FlrSessionStatus": "MountCompletedSuccessfully"})
    assert is_live_session({"FlrSessionStatus": "MountInProgress"})
    # unmounted/ended/failed sessions linger as history and hold nothing
    assert not is_live_session({"FlrSessionStatus": "UnmountCompletedSuccessfully"})
    assert not is_live_session({"FlrSessionStatus": "MountFailed"})
    assert not is_live_session({})


class _FakeClient:
    """Stands in for ZertoClient in teardown tests."""

    def __init__(self, sessions=None, fail_end=False):
        self._sessions = sessions or []
        self.fail_end = fail_end
        self.ended: list[str] = []

    async def list_flrs(self):
        return self._sessions

    async def end_flr(self, session_id):
        from zerto_rewind_mcp.client import ZertoError

        if self.fail_end:
            raise ZertoError("unmount refused")
        self.ended.append(session_id)


def test_teardown_reports_failure_instead_of_swallowing():
    import asyncio

    from zerto_rewind_mcp.server import _teardown_flr

    client = _FakeClient(fail_end=True)
    out = asyncio.run(_teardown_flr(client, "sess-1", set()))
    assert out["ok"] is False
    assert out["failed"][0]["session_id"] == "sess-1"
    assert "still be mounted" in out["message"]


def test_teardown_reaps_orphan_when_session_id_never_parsed():
    import asyncio

    from zerto_rewind_mcp.server import _teardown_flr

    # start_flr succeeded on the ZVM but the response did not parse, so the
    # caller never learned the id. The new live session must still be reaped.
    live = [{"FlrSessionIdentifier": "new-1", "FlrSessionStatus": "MountCompletedSuccessfully"}]
    client = _FakeClient(sessions=live)
    out = asyncio.run(_teardown_flr(client, None, set()))
    assert out["ended"] == ["new-1"]
    assert out["orphans_reaped"] == ["new-1"]
    assert client.ended == ["new-1"]


def test_teardown_leaves_pre_existing_sessions_alone():
    import asyncio

    from zerto_rewind_mcp.server import _teardown_flr

    live = [
        {"FlrSessionIdentifier": "someone-else", "FlrSessionStatus": "MountCompletedSuccessfully"}
    ]
    client = _FakeClient(sessions=live)
    out = asyncio.run(_teardown_flr(client, None, {"someone-else"}))
    assert out["ended"] == []
    assert client.ended == []