folding-studio/tests/commands/test_experiment.py

from pathlib import Path
from unittest import mock

import pytest
from folding_studio.cli import app
from folding_studio.config import API_URL, REQUEST_TIMEOUT
from typer.testing import CliRunner
from folding_studio.utils.headers import get_auth_headers

current_workdir = Path(__file__).parent.resolve()
data_dir = Path(current_workdir / "data")

runner = CliRunner()


@pytest.fixture(autouse=True)
def mock_get_auth_headers(request):
    if "apikeytest" in request.keywords:
        return_value = get_auth_headers(api_key="MY_KEY")
    else:
        return_value = {"Authorization": "Bearer identity_token"}
    with mock.patch(
        "folding_studio.commands.experiment.get_auth_headers", return_value=return_value
    ) as m:
        yield m



@pytest.fixture()
def mock_request_get():
    with mock.patch("requests.get") as m:
        yield m


def test_experiment_status_pass(mock_request_get: mock.Mock, headers):
    mock_response = mock.MagicMock()
    mock_response.ok = True
    mock_response.json.return_value = {"status": "Done."}
    mock_request_get.return_value = mock_response

    result = runner.invoke(app, ["experiment", "status", "ID"])
    assert result.exit_code == 0
    assert "Done." in result.stdout
    mock_request_get.assert_called_once_with(
        API_URL + "getExperimentStatus",
        params={"experiment_id": "ID"},
        headers=headers,
        timeout=REQUEST_TIMEOUT,
    )


def test_experiment_list_pass(mock_request_get: mock.Mock, tmp_files, headers):
    mock_response = mock.MagicMock()
    mock_response.ok = True
    mock_response.json.return_value = {
        "Done": ["exp_id_1", "exp_id_2"],
        "Pending": ["exp_id_3"],
    }
    mock_request_get.return_value = mock_response

    result = runner.invoke(
        app, ["experiment", "list", "--output", str(tmp_files["metadata_file"])]
    )
    assert result.exit_code == 0
    mock_request_get.assert_called_once_with(
        API_URL + "getDoneAndPendingExperiments",
        headers=headers,
        timeout=REQUEST_TIMEOUT,
    )