open-metadata--openmetadata/ingestion/tests/unit/sdk/test_restore_async.py

#  Copyright 2026 Collate
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#  http://www.apache.org/licenses/LICENSE-2.0
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
"""Tests for the fluent restore + server-side async option (issue #4003)."""

from unittest.mock import MagicMock

import pytest

from metadata.sdk import Tables
from metadata.sdk.entities.base import AsyncJobResponse, RestoreOperation


@pytest.fixture
def mock_client():
    client = MagicMock()
    client.get_suffix.return_value = "/tables"
    rest_client = MagicMock()
    client.client = rest_client
    Tables.set_default_client(client)
    return client


def _table_payload(table_id: str) -> dict:
    """Minimum dict shape that pydantic_core accepts as a Table."""
    return {
        "id": table_id,
        "name": "t",
        "fullyQualifiedName": "service.db.schema.t",
        "deleted": False,
        "columns": [],
    }


def test_restore_sync_calls_put_without_async_param(mock_client):
    table_id = "b67eac63-9e43-41f5-afb9-387c85df1d8b"
    rest_client = mock_client.client
    rest_client.put.return_value = _table_payload(table_id)

    Tables.restore(table_id)

    rest_client.put.assert_called_once()
    path = rest_client.put.call_args[0][0]
    assert path.endswith("/restore")
    assert "async=true" not in path


def test_restore_async_appends_async_query_param(mock_client):
    table_id = "b67eac63-9e43-41f5-afb9-387c85df1d8b"
    rest_client = mock_client.client
    rest_client.put.return_value = {"jobId": "job-42", "message": "Restore initiated"}

    response = Tables.restore_async(table_id)

    assert isinstance(response, AsyncJobResponse)
    assert response.job_id == "job-42"
    assert response.message == "Restore initiated"
    rest_client.put.assert_called_once()
    path = rest_client.put.call_args[0][0]
    assert path.endswith("/restore?async=true")


def test_fluent_restore_request_sync_returns_entity(mock_client):
    table_id = "b67eac63-9e43-41f5-afb9-387c85df1d8b"
    rest_client = mock_client.client
    rest_client.put.return_value = _table_payload(table_id)

    op = Tables.restore_request(table_id)
    assert isinstance(op, RestoreOperation)
    op.execute()

    path = rest_client.put.call_args[0][0]
    assert "async=true" not in path


def test_fluent_restore_request_with_async_returns_job_response(mock_client):
    table_id = "b67eac63-9e43-41f5-afb9-387c85df1d8b"
    rest_client = mock_client.client
    rest_client.put.return_value = {"jobId": "job-7", "message": "Restore initiated"}

    job = Tables.restore_request(table_id).with_async().execute()

    assert isinstance(job, AsyncJobResponse)
    assert job.job_id == "job-7"
    assert "async=true" in rest_client.put.call_args[0][0]


def test_async_job_response_from_response_handles_dict():
    response = AsyncJobResponse.from_response({"jobId": "abc", "message": "ok"})
    assert response.job_id == "abc"
    assert response.message == "ok"


def test_async_job_response_from_response_passes_through_existing():
    original = AsyncJobResponse(job_id="abc", message="ok")
    assert AsyncJobResponse.from_response(original) is original


def test_async_job_response_from_response_rejects_unknown_type():
    with pytest.raises(TypeError):
        AsyncJobResponse.from_response("not a dict")


def test_async_job_response_from_response_rejects_missing_job_id():
    with pytest.raises(ValueError, match="non-empty jobId"):
        AsyncJobResponse.from_response({"message": "no id here"})


def test_async_job_response_from_response_rejects_empty_job_id():
    with pytest.raises(ValueError, match="non-empty jobId"):
        AsyncJobResponse.from_response({"jobId": "", "message": "blank"})