class TestRedisSensor(unittest.TestCase):

    def setUp(self):
        configuration.load_test_config()
        args = {
            'owner': 'airflow',
            'start_date': DEFAULT_DATE
        }

        self.dag = DAG('test_dag_id', default_args=args)
        self.sensor = RedisKeySensor(
            task_id='test_task',
            redis_conn_id='redis_default',
            dag=self.dag,
            key='test_key'
        )

    @patch("airflow.contrib.hooks.redis_hook.RedisHook.key_exists")
    def test_poke(self, key_exists):
        key_exists.return_value = True
        self.assertTrue(self.sensor.poke(None))

        key_exists.return_value = False
        self.assertFalse(self.sensor.poke(None))

    @patch("airflow.contrib.hooks.redis_hook.StrictRedis.exists")
    def test_existing_key_called(self, redis_client_exists):
        self.sensor.run(
            start_date=DEFAULT_DATE,
            end_date=DEFAULT_DATE, ignore_ti_state=True
        )

        self.assertTrue(redis_client_exists.called_with('test_key'))
예제 #2
0
class TestRedisSensor(unittest.TestCase):
    def setUp(self):
        configuration.load_test_config()
        args = {'owner': 'airflow', 'start_date': DEFAULT_DATE}

        self.dag = DAG('test_dag_id', default_args=args)
        self.sensor = RedisKeySensor(task_id='test_task',
                                     redis_conn_id='redis_default',
                                     dag=self.dag,
                                     key='test_key')

    @patch("airflow.contrib.hooks.redis_hook.RedisHook.key_exists")
    def test_poke(self, key_exists):
        key_exists.return_value = True
        self.assertTrue(self.sensor.poke(None))

        key_exists.return_value = False
        self.assertFalse(self.sensor.poke(None))

    @patch("airflow.contrib.hooks.redis_hook.StrictRedis.exists")
    def test_existing_key_called(self, redis_client_exists):
        self.sensor.run(start_date=DEFAULT_DATE,
                        end_date=DEFAULT_DATE,
                        ignore_ti_state=True)

        self.assertTrue(redis_client_exists.called_with('test_key'))
예제 #3
0
    def setUp(self):
        args = {'owner': 'airflow', 'start_date': DEFAULT_DATE}

        self.dag = DAG('test_dag_id', default_args=args)
        self.sensor = RedisKeySensor(task_id='test_task',
                                     redis_conn_id='redis_default',
                                     dag=self.dag,
                                     key='test_key')
예제 #4
0
class TestRedisSensor(unittest.TestCase):
    def setUp(self):
        args = {'owner': 'airflow', 'start_date': DEFAULT_DATE}

        self.dag = DAG('test_dag_id', default_args=args)
        self.sensor = RedisKeySensor(task_id='test_task',
                                     redis_conn_id='redis_default',
                                     dag=self.dag,
                                     key='test_key')

    def test_poke(self):
        hook = RedisHook(redis_conn_id='redis_default')
        redis = hook.get_conn()
        redis.set('test_key', 'test_value')
        self.assertTrue(self.sensor.poke(None), "Key exists on first call.")
        redis.delete('test_key')
        self.assertFalse(self.sensor.poke(None),
                         "Key does NOT exists on second call.")
    def setUp(self):
        configuration.load_test_config()
        args = {
            'owner': 'airflow',
            'start_date': DEFAULT_DATE
        }

        self.dag = DAG('test_dag_id', default_args=args)
        self.sensor = RedisKeySensor(
            task_id='test_task',
            redis_conn_id='redis_default',
            dag=self.dag,
            key='test_key'
        )
예제 #6
0
class TestRedisSensor(unittest.TestCase):

    def setUp(self):
        configuration.load_test_config()
        args = {
            'owner': 'airflow',
            'start_date': DEFAULT_DATE
        }

        self.dag = DAG('test_dag_id', default_args=args)
        self.sensor = RedisKeySensor(
            task_id='test_task',
            redis_conn_id='redis_default',
            dag=self.dag,
            key='test_key'
        )

    def test_poke(self):
        hook = RedisHook(redis_conn_id='redis_default')
        redis = hook.get_conn()
        redis.set('test_key', 'test_value')
        self.assertTrue(self.sensor.poke(None), "Key exists on first call.")
        redis.delete('test_key')
        self.assertFalse(self.sensor.poke(None), "Key does NOT exists on second call.")
# from airflow.contrib.sensors import RedisKeySensor
import redis

from airflow.contrib.hooks.redis_hook import RedisHook
from airflow.sensors.base_sensor_operator import BaseSensorOperator
from airflow.utils.decorators import apply_defaults

args = {
    'start_date': datetime.utcnow(),
    'owner': 'airflow',
}

dag = DAG(dag_id='vlass_execute', default_args=args, schedule_interval=None)

sensor = RedisKeySensor(task_id='check_task',
                        redis_conn_id='redis_default',
                        dag=dag,
                        key='test_key')

start_task = DummyOperator(task_id='start_task', dag=dag)
end_task = DummyOperator(task_id='end_task', dag=dag)


def get_file_names():
    redis_conn = redis.StrictRedis()
    results = redis_conn.get('test_key')
    if results is not None:
        file_name_list = results.decode('utf-8').split()[1:]
        # redis_conn.delete('test_key')
        return file_name_list
    else:
        return []