import unittest
from zenyatta.common.util import get_task_instance
from zenyatta.common import xcoms
from zenyatta.db.sql import run_table_update

var = {
        'airflow_conn': {
            'host': 'zenyatta-postgres.cifgffw7w2ar.us-west-2.rds.amazonaws.com',
            'schema': 'airflow',
            'login': 'airflow',
            'password': 'airflow1234'},
        'task_id': 'dummy-task-id',
        'dag_id': 'dummy-dag',
        'execution_date': '2017-07-01 00:00:00',
        'ts_nodash': '20170101T000000',
        'volume_id': 'vol-12345',
        'rds_id': 'rds-12345',
        'instance_id': 'ins-12345',
        'conn_id': 'conn-12345',
        'pg_isready': 'readyornot',
        'spark_script': 'whereismyspark',
        'spark_config': 'whataconf',
        'spark_submit_flag': 'False',
        'db_tmp_conn_id': 'db-conn-12345',
        'csv': 'ohmycsvfile',
        'bucket': 'ohmybucket',
        'role_arn': 'dummy_arn',
        's3_script_path': 's3file.csv',
        'local_script': 'whereismylocal',
        's3_parquet': 'farawayons3',
        'local_parquet': 'localonmyebs'
    }


class TaskInstance:
    def __init__(self, dag_id, ts_nodash, task_id):
        self.dag_id = dag_id
        self.ts_nodash = ts_nodash
        self.task_id = task_id


class XcomTestCase(unittest.TestCase):

    # create a dummy task instance to play with
    @unittest.skip("pass and skip")
    def get_dummy_task_instance(self):
        stmt = "insert into task_instance (task_id, dag_id, execution_date) " \
               "values ('{task_id}', '{dag_id}', '{execution_date}');" \
               .format(task_id=var['task_id'], dag_id=var['dag_id'], execution_date=var['execution_date'])
        try:
            rtn = run_table_update(var['airflow_conn']['host'],
                                   var['airflow_conn']['schema'],
                                   var['airflow_conn']['login'],
                                   var['airflow_conn']['password'],
                                   stmt)
        except Exception:
            pass

    def test_volume_id_xcoms(self):
        task_instance = TaskInstance(var['dag_id'], var['execution_date'], var['task_id'])
        key = xcoms.XcomVolumeIdKey(task_instance, var['ts_nodash']).get_key()
        self.assertEqual(key, var['dag_id'] + '-volume_id-' + var['ts_nodash'])
        # task_instance.xcom_push(key=key, value=var['volume_id'])
        # value = task_instance.xcom_pull(key=key, task_ids=var['task_id'])
        # self.assertEqual(value, var['volume_id'])

    def test_instance_id_xcoms(self):
        task_instance = TaskInstance(var['dag_id'], var['execution_date'], var['task_id'])
        key = xcoms.XcomInstanceIdKey(task_instance, var['ts_nodash']).get_key()
        self.assertEqual(key, var['dag_id'] + '-instance_id-' + var['ts_nodash'])
        # task_instance.xcom_push(key=key, value=var['instance_id'])
        # value = task_instance.xcom_pull(key=key, task_ids=var['task_id'])
        # self.assertEqual(value, var['instance_id'])

    def test_conn_id_xcoms(self):
        task_instance = TaskInstance(var['dag_id'], var['execution_date'], var['task_id'])
        key = xcoms.XcomConnIdKey(task_instance, var['ts_nodash']).get_key()
        self.assertEqual(key, var['dag_id'] + '-conn_id-' + var['ts_nodash'])
        # task_instance.xcom_push(key=key, value=var['conn_id'])
        # value = task_instance.xcom_pull(key=key, task_ids=var['task_id'])
        # self.assertEqual(value, var['conn_id'])

    def test_pgread_id_xcoms(self):
        task_instance = TaskInstance(var['dag_id'], var['execution_date'], var['task_id'])
        key = xcoms.XcomPGReadyKey(task_instance, var['ts_nodash']).get_key()
        self.assertEqual(key, var['dag_id'] + '-pg_isready-' + var['ts_nodash'])
        # task_instance.xcom_push(key=key, value=var['pg_isready'])
        # value = task_instance.xcom_pull(key=key, task_ids=var['task_id'])
        # self.assertEqual(value, var['pg_isready'])

    def test_spark_script_xcoms(self):
        task_instance = TaskInstance(var['dag_id'], var['execution_date'], var['task_id'])
        key = xcoms.XcomSparkScriptKey(task_instance, var['ts_nodash']).get_key()
        self.assertEqual(key, var['dag_id'] + '-spark_script-' + var['ts_nodash'])
        # task_instance.xcom_push(key=key, value=var['spark_script'])
        # value = task_instance.xcom_pull(key=key, task_ids=var['task_id'])
        # self.assertEqual(value, var['spark_script'])

    def test_spark_conf_xcoms(self):
        task_instance = TaskInstance(var['dag_id'], var['execution_date'], var['task_id'])
        key = xcoms.XcomSparkConfigKey(task_instance, var['ts_nodash']).get_key()
        self.assertEqual(key, var['dag_id'] + '-spark_config-' + var['ts_nodash'])
        # task_instance.xcom_push(key=key, value=var['spark_config'])
        # value = task_instance.xcom_pull(key=key, task_ids=var['task_id'])
        # self.assertEqual(value, var['spark_config'])

    def test_spark_submit_xcoms(self):
        task_instance = TaskInstance(var['dag_id'], var['execution_date'], var['task_id'])
        key = xcoms.XcomSparkSubmitKey(task_instance, var['ts_nodash']).get_key()
        self.assertEqual(key, var['dag_id'] + '-spark_submit_flag-' + var['ts_nodash'])
        # task_instance.xcom_push(key=key, value=var['spark_submit_flag'])
        # value = task_instance.xcom_pull(key=key, task_ids=var['task_id'])
        # self.assertEqual(value, var['spark_submit_flag'])

    def test_rds_id_xcoms(self):
        task_instance = TaskInstance(var['dag_id'], var['execution_date'], var['task_id'])
        key = xcoms.XcomRDSIdKey(task_instance, var['ts_nodash']).get_key()
        self.assertEqual(key, var['dag_id'] + '-rds_id-' + var['ts_nodash'])
        # task_instance.xcom_push(key=key, value=var['rds_id'])
        # value = task_instance.xcom_pull(key=key, task_ids=var['task_id'])
        # self.assertEqual(value, var['rds_id'])

    def test_db_conn_id_xcoms(self):
        task_instance = TaskInstance(var['dag_id'], var['execution_date'], var['task_id'])
        key = xcoms.XcomDBConnIdKey(task_instance, var['ts_nodash']).get_key()
        self.assertEqual(key, var['dag_id'] + '-db_tmp_conn_id-' + var['ts_nodash'])
        # task_instance.xcom_push(key=key, value=var['db_tmp_conn_id'])
        # value = task_instance.xcom_pull(key=key, task_ids=var['task_id'])
        # self.assertEqual(value, var['db_tmp_conn_id'])

    def test_csv_xcoms(self):
        task_instance = TaskInstance(var['dag_id'], var['execution_date'], var['task_id'])
        key = xcoms.XcomCSVKey(task_instance, var['ts_nodash']).get_key()
        self.assertEqual(key, var['dag_id'] + '-csv-' + var['ts_nodash'])
        # task_instance.xcom_push(key=key, value=var['csv'])
        # value = task_instance.xcom_pull(key=key, task_ids=var['task_id'])
        # self.assertEqual(value, var['csv'])

    def test_bucket_xcoms(self):
        task_instance = TaskInstance(var['dag_id'], var['execution_date'], var['task_id'])
        key = xcoms.XcomBucketKey(task_instance, var['ts_nodash']).get_key()
        self.assertEqual(key, var['dag_id'] + '-bucket-' + var['ts_nodash'])
        # task_instance.xcom_push(key=key, value=var['bucket'])
        # value = task_instance.xcom_pull(key=key, task_ids=var['task_id'])
        # self.assertEqual(value, var['bucket'])

    def test_role_arn_xcoms(self):
        task_instance = TaskInstance(var['dag_id'], var['execution_date'], var['task_id'])
        key = xcoms.XcomRoleArnKey(task_instance, var['ts_nodash']).get_key()
        self.assertEqual(key, var['dag_id'] + '-role_arn-' + var['ts_nodash'])
        # task_instance.xcom_push(key=key, value=var['role_arn'])
        # value = task_instance.xcom_pull(key=key, task_ids=var['task_id'])
        # self.assertEqual(value, var['role_arn'])

    def test_s3_script_xcoms(self):
        task_instance = TaskInstance(var['dag_id'], var['execution_date'], var['task_id'])
        key = xcoms.XcomS3ScriptKey(task_instance, var['ts_nodash']).get_key()
        self.assertEqual(key, var['dag_id'] + '-s3_script_path-' + var['ts_nodash'])
        # task_instance.xcom_push(key=key, value=var['s3_script_path'])
        # value = task_instance.xcom_pull(key=key, task_ids=var['task_id'])
        # self.assertEqual(value, var['s3_script_path'])

    def test_local_script_xcoms(self):
        task_instance = TaskInstance(var['dag_id'], var['execution_date'], var['task_id'])
        key = xcoms.XcomLocalScriptKey(task_instance, var['ts_nodash']).get_key()
        self.assertEqual(key, var['dag_id'] + '-local_script-' + var['ts_nodash'])
        # task_instance.xcom_push(key=key, value=var['local_script'])
        # value = task_instance.xcom_pull(key=key, task_ids=var['task_id'])
        # self.assertEqual(value, var['local_script'])

    def test_s3_pq_xcoms(self):
        task_instance = TaskInstance(var['dag_id'], var['execution_date'], var['task_id'])
        key = xcoms.XcomS3PQKey(task_instance, var['ts_nodash']).get_key()
        self.assertEqual(key, var['dag_id'] + '-s3_parquet-' + var['ts_nodash'])
        # task_instance.xcom_push(key=key, value=var['s3_parquet'])
        # value = task_instance.xcom_pull(key=key, task_ids=var['task_id'])
        # self.assertEqual(value, var['s3_parquet'])

    def test_local_pq_xcoms(self):
        task_instance = TaskInstance(var['dag_id'], var['execution_date'], var['task_id'])
        key = xcoms.XcomLocalPQKey(task_instance, var['ts_nodash']).get_key()
        self.assertEqual(key, var['dag_id'] + '-local_parquet-' + var['ts_nodash'])
        # task_instance.xcom_push(key=key, value=var['local_parquet'])
        # value = task_instance.xcom_pull(key=key, task_ids=var['task_id'])
        # self.assertEqual(value, var['local_parquet'])


if __name__ == '__main__':
    unittest.main()
