from unittest.mock import AsyncMock, call, patch, ANY
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch


@fixture
def predictions_batch() -> PredictionsBatch:
    return PredictionsBatch()


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


@mark.asyncio
@patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
    workflow_mock.execute_local_activity_method.return_value = {
        'data': 'test_data'
    }
    input_data = {
        'schedule_name': 'test_schedule',
        'model_name': 'test_model',
        'model_id': 'test_model_id',
        'query': 'SELECT * FROM test',
        'schema': 'test_schema',
        'table_name': 'test_table',
        'opc_output_config': 'test_opc_output_config',
        'datetime_columns': ['timestamp', 'created_at'],
        'prediction_store_policy': 'erl:1',
        'model_config': {
            'retention': '30'
        }
    }

    await predictions_batch.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
        )
    ])
    prediction_input = {
        'metadata': metadata,
        'data': {'data': 'test_data'},
        'schema': input_data['schema'],
        'table_name': input_data['table_name'],
        'model_id': input_data['model_id'],
        'model_name': input_data['model_name'],
        'input_filters': input_data.get('input_filters', {
            'EMPTY_DATA': {
                'POLICY': 'STOP'
            }
        }),
        'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {
            'API_ERROR': {
                'POLICY': 'STOP'
            }
        }),
        'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {
            'API_ERROR': {
                'POLICY': 'STOP'
            }
        }),
        'model_config': input_data.get('model_config', {}),
        'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
        'opc_output_config': input_data.get('opc_output_config', {}),
        'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1')
    }

    workflow_mock.execute_child_workflow.assert_has_calls([
        call(
            'prediction_process', prediction_input)
    ])