import unittest
from unittest import mock

from parameterized import parameterized

from exp_dbx_runner_lambda.entities import Config, Input
from exp_dbx_runner_lambda.dbx_service import DatabricksJob, DatabricksService, RunState
from exp_dbx_runner_lambda.exceptions import DbUpdateError


class DatabricksServiceTestCase(unittest.TestCase):
    def setUp(self) -> None:
        self.config = Config(
            environment="dev", dbx_url="dbx_url", dbx_token="dbx_token", use_payload_buckets=False,
            sentry_secret_key='test_key'
        )
        self.payload = Input(
            source_path="s3://decompressed/gras/file.txt",
            uow_id="gras-20190911-sme-dim_product_track-v1",
            disassemble_content_status_id=123,
            sf_token="test_token",
            sf_execution_name="test_name",
            databricks_job_id=1,
            databricks_spark_params=[
                "-Dconfig.resource=application.dev.conf", "--drop-from-head", "3",
                "--drop-from-tail", "2"
            ],
            dest_bucket="dest"
        )
        self.pg_repo = mock.Mock()
        self.dbx_service = DatabricksService(
            mock.Mock(), mock.Mock(), self.config, self.payload, self.pg_repo
        )

    @parameterized.expand(
        [
            (
                False, [
                    "-Dconfig.resource=application.dev.conf", "--drop-from-head", "3",
                    "--drop-from-tail", "2", "--sf-token", "test_token", "--sf-name", "test_name",
                    "--source-key", "gras/file.txt", "--disassemble-content-status-id", "123"
                ]
            ),
            (
                True, [
                    "-Dconfig.resource=application.dev.conf -Ddelphi.etl.exploration.input-bucket=decompressed -Ddelphi.etl.exploration.output-bucket=dest",
                    "--drop-from-head", "3", "--drop-from-tail", "2", "--sf-token", "test_token",
                    "--sf-name", "test_name", "--source-key", "gras/file.txt",
                    "--disassemble-content-status-id", "123"
                ]
            ),
        ]
    )
    def test_prepare_dbx_command(self, use_buckets, expected):
        self.config.use_payload_buckets = use_buckets
        actual = self.dbx_service._prepare_dbx_command()
        self.assertEqual(actual, expected)

    def test_update_db_ok(self):
        self.pg_repo.update_dbx_run_id.side_effect = lambda *args: True
        result = self.dbx_service._update_pg_record(123)
        self.assertEqual(result, True)

    def test_update_db_correct_raises_exc(self):
        self.pg_repo.update_dbx_run_id.side_effect = lambda *args: False
        self.assertRaises(DbUpdateError, self.dbx_service._update_pg_record, 123)


class DatabricksJobTestCase(unittest.TestCase):
    def setUp(self) -> None:
        self.dbx_service = DatabricksJob(
            logger=mock.Mock(), dbx_url="dbx_url", dbx_token="dbx_token"
        )

    @unittest.mock.patch('requests.post')
    def test_start_job(self, mocked_post):
        run_id, job_id = 10, 12
        mocked_post.return_value.json.return_value = {'run_id': run_id}
        command = ["--test1", "--test2"]
        actual_run_id = self.dbx_service.start_job(job_id, command)
        mocked_post.assert_called_once_with(
            url='https://dbx_url/api/2.0/jobs/run-now',
            headers={'Authorization': 'Bearer dbx_token'},
            json={
                'job_id': job_id,
                'spark_submit_params': command
            },
        )
        self.assertEqual(actual_run_id, run_id)

    @unittest.mock.patch('requests.get')
    def test_get_run_state(self, mocked_get):
        mocked_get.return_value.json.return_value = {'state': {"life_cycle_state": 'IN_PROGRESS'}}
        run_id = 10
        _state = self.dbx_service.get_run_state(run_id)
        self.assertEqual(_state, RunState(life_cycle_state='IN_PROGRESS'))
        mocked_get.assert_called_once_with(
            url='https://dbx_url/api/2.0/jobs/runs/get',
            headers={'Authorization': 'Bearer dbx_token'},
            params={'run_id': run_id},
        )