import unittest.mock

import pytest
from db_schema.factories.apps import DatabricksExecutionFactory
from db_schema.schemas.apps import DBXJobStatusEnum

from apps_etl_databricks.objects.run import DBXRuns, PersistentDBXRun, RunState


def test_run_state_running():
    state = RunState(life_cycle_state='RUNNING', )
    assert not state.is_completed()
    assert not state.is_failed()
    assert not state.is_skipped()
    assert not state.is_cancelled()
    assert state.is_running()


def test_run_state_skipped():
    state = RunState(life_cycle_state='SKIPPED', )
    assert not state.is_completed()
    assert not state.is_failed()
    assert not state.is_cancelled()
    assert not state.is_running()
    assert state.is_skipped()


def test_run_state_completed():
    state = RunState(
        life_cycle_state='TERMINATED',
        result_state='SUCCESS',
    )
    assert not state.is_failed()
    assert not state.is_skipped()
    assert not state.is_cancelled()
    assert not state.is_running()
    assert state.is_completed()


def test_run_state_failed():
    state = RunState(
        life_cycle_state='TERMINATED',
        result_state='FAILED',
    )
    assert not state.is_completed()
    assert not state.is_skipped()
    assert not state.is_cancelled()
    assert not state.is_running()
    assert state.is_failed()


def test_run_state_cancelled():
    state = RunState(
        life_cycle_state='TERMINATED',
        result_state='CANCELED',
    )
    assert not state.is_completed()
    assert not state.is_failed()
    assert not state.is_skipped()
    assert not state.is_running()
    assert state.is_cancelled()


@unittest.mock.patch('requests.get')
def test_dbx_runs_state(mocked_get):
    expected_result = RunState(life_cycle_state='RUNNING')

    mocked_get.return_value.json.return_value = {'state': expected_result}

    runs = DBXRuns(unittest.mock.Mock(), 'test-dbx-url.com', 'test-token')
    run_state = runs.state(10)

    assert run_state == expected_result


@pytest.mark.integration
def test_persistent_dbx_run(session):
    dbx_execution = DatabricksExecutionFactory()

    dbx_runs = unittest.mock.Mock()
    dbx_runs.state.return_value = RunState(
        life_cycle_state='TERMINATED',
        result_state='FAILED',
    )

    persistent_dbx_run = PersistentDBXRun(
        logger=unittest.mock.Mock(),
        runs=dbx_runs,
        session=session,
        sf_name=dbx_execution.sf_execution_name,
        sqs_service=unittest.mock.Mock(),
        conversion_params={},
    )

    assert persistent_dbx_run.sync()
    assert dbx_execution.completed_at


@pytest.mark.integration
def test_persistent_dbx_run_cancelled(session):
    dbx_execution = DatabricksExecutionFactory()

    dbx_runs = unittest.mock.Mock()
    dbx_runs.state.return_value = RunState(
        life_cycle_state='TERMINATED',
        result_state='CANCELED',
    )

    persistent_dbx_run = PersistentDBXRun(
        logger=unittest.mock.Mock(),
        runs=dbx_runs,
        session=session,
        sf_name=dbx_execution.sf_execution_name,
        sqs_service=unittest.mock.Mock(),
        conversion_params={},
    )

    assert persistent_dbx_run.sync()
    assert dbx_execution.status == DBXJobStatusEnum.CANCELLED
    assert dbx_execution.completed_at


@pytest.mark.integration
def test_run_send_conversion_metadata(session):
    dbx_execution = DatabricksExecutionFactory()
    dbx_runs = unittest.mock.Mock()
    sqs_service = unittest.mock.Mock()
    sqs_service.push.return_value = True

    persistent_dbx_run = PersistentDBXRun(
        logger=unittest.mock.Mock(),
        runs=dbx_runs,
        session=session,
        sf_name=dbx_execution.sf_execution_name,
        sqs_service=sqs_service,
        conversion_params={'contexts': []},
    )

    assert persistent_dbx_run.send_conversion_metadata(dbx_execution.unit_of_work_id)