from unittest.mock import AsyncMock, MagicMock, call, patch, ANY
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.minimal_retrain import MinimalRetrain


@fixture
def minimal_retrain() -> MinimalRetrain:
    return MinimalRetrain()


metadata = {
    "metadata": {
        "model_id": "test_model_id",
        "model_name": "test_model",
        "workflow_name": "minimal_retrain",
        "schedule_name": "test_schedule",
    },
}


@mark.asyncio
@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
    input_data = {
        "model_id": "test_model_id",
        "model_name": "test_model",
        "workflow_name": "minimal_retrain",
        "schedule_name": "test_schedule",
        "query": "test_query",
        "schema": "test_schema",
        "table_name": "test_table",
    }

    workflow_mock.execute_activity_method = AsyncMock(
        return_value={
            "data1": "1",
            "data2": "2",
        }
    )

    await minimal_retrain.run(input_data)

    workflow_mock.execute_local_activity_method.assert_has_calls(
        [
            call(
                Activities.load_custom_query,
                {
                    **metadata,
                    "query": input_data["query"],
                    'datetime_columns': input_data.get('datetime_columns', [])
                },
                retry_policy=ANY,
                start_to_close_timeout=ANY
            )
        ]
    )

    workflow_mock.execute_activity_method.assert_has_calls([
        call(
            Activities.retrain_model,
            {
                **metadata,
                'data': workflow_mock.execute_local_activity_method.return_value,
                'model_name': input_data['model_name'],
            },
            retry_policy=ANY,
            start_to_close_timeout=ANY
        )
    ])

    workflow_mock.execute_activity_method.assert_has_calls([
        call(
            Activities.update_production_model,
            {
                **metadata,
                'model_name': input_data['model_name'],
                'model_id': input_data['model_id'],
                **workflow_mock.execute_activity_method.return_value,
            },
            retry_policy=ANY,
            start_to_close_timeout=ANY
        )
    ])

    workflow_mock.execute_activity_method.assert_has_calls([
        call(
            Activities.export_data_to_postgres,
            {
                **metadata,
                'data': workflow_mock.execute_activity_method.return_value,
                'schema': input_data['schema'],
                'table_name': input_data['table_name'],
            },
            retry_policy=ANY,
            start_to_close_timeout=ANY
        )
    ])