def test_poke_true(self):
sensor = RedisPubSubSensor(task_id='test_task',
dag=self.dag,
channels='test',
redis_conn_id='redis_default')
hook = RedisHook(redis_conn_id='redis_default')
redis = hook.get_conn()
redis.publish('test', 'message')
result = sensor.poke(self.mock_context)
self.assertFalse(result)
result = sensor.poke(self.mock_context)
self.assertTrue(result)
context_calls = [
call.xcom_push(key='message',
value={
'type': 'message',
'pattern': None,
'channel': b'test',
'data': b'message'
})
]
self.assertTrue(self.mock_context['ti'].method_calls == context_calls,
"context calls should be same")
result = sensor.poke(self.mock_context)
self.assertFalse(result)
def test_poke_mock_true(self, mock_redis_conn):
sensor = RedisPubSubSensor(task_id='test_task',
dag=self.dag,
channels='test',
redis_conn_id='redis_default')
mock_redis_conn().pubsub().get_message.return_value = {
'type': 'message',
'channel': b'test',
'data': b'd1',
}
result = sensor.poke(self.mock_context)
self.assertTrue(result)
context_calls = [
call.xcom_push(key='message',
value={
'type': 'message',
'channel': b'test',
'data': b'd1'
})
]
self.assertTrue(self.mock_context['ti'].method_calls == context_calls,
"context call should be same")
def test_poke_false(self):
sensor = RedisPubSubSensor(task_id='test_task',
dag=self.dag,
channels='test',
redis_conn_id='redis_default')
result = sensor.poke(self.mock_context)
self.assertFalse(result)
self.assertTrue(self.mock_context['ti'].method_calls == [],
"context calls should be same")
result = sensor.poke(self.mock_context)
self.assertFalse(result)
self.assertTrue(self.mock_context['ti'].method_calls == [],
"context calls should be same")
def test_poke_mock_false(self, mock_redis_conn):
sensor = RedisPubSubSensor(task_id='test_task',
dag=self.dag,
channels='test',
redis_conn_id='redis_default')
mock_redis_conn().pubsub().get_message.return_value = {
'type': 'subscribe',
'channel': b'test',
'data': b'd1',
}
result = sensor.poke(self.mock_context)
assert not result
context_calls = []
assert self.mock_context[
'ti'].method_calls == context_calls, "context calls should be same"