def test_read_with_empty_metadata(self):
ts = pendulum.now()
logs, metadatas = self.es_task_handler.read(self.ti, 1, {})
self.assertEqual(1, len(logs))
self.assertEqual(len(logs), len(metadatas))
self.assertEqual(self.test_message, logs[0])
self.assertFalse(metadatas[0]['end_of_log'])
# offset should be initialized to 0 if not provided.
self.assertEqual(1, metadatas[0]['offset'])
# last_log_timestamp will be initialized using log reading time
# if not last_log_timestamp is provided.
self.assertTrue(
timezone.parse(metadatas[0]['last_log_timestamp']) > ts)
# case where offset is missing but metadata not empty.
self.es.delete(index=self.index_name, doc_type=self.doc_type, id=1)
logs, metadatas = self.es_task_handler.read(self.ti, 1,
{'end_of_log': False})
self.assertEqual(1, len(logs))
self.assertEqual(len(logs), len(metadatas))
self.assertEqual([''], logs)
self.assertFalse(metadatas[0]['end_of_log'])
# offset should be initialized to 0 if not provided.
self.assertEqual(0, metadatas[0]['offset'])
# last_log_timestamp will be initialized using log reading time
# if not last_log_timestamp is provided.
self.assertTrue(
timezone.parse(metadatas[0]['last_log_timestamp']) > ts)
def test_read_with_none_meatadata(self):
logs, metadatas = self.es_task_handler.read(self.ti, 1)
self.assertEqual(1, len(logs))
self.assertEqual(len(logs), len(metadatas))
self.assertEqual(self.test_message, logs[0])
self.assertFalse(metadatas[0]['end_of_log'])
self.assertEqual(1, metadatas[0]['offset'])
self.assertTrue(
timezone.parse(metadatas[0]['last_log_timestamp']) <
pendulum.now())
def test_read(self):
ts = pendulum.now()
logs, metadatas = self.es_task_handler.read(
self.ti, 1, {
'offset': 0,
'last_log_timestamp': str(ts),
'end_of_log': False
})
self.assertEqual(1, len(logs))
self.assertEqual(len(logs), len(metadatas))
self.assertEqual(self.test_message, logs[0])
self.assertFalse(metadatas[0]['end_of_log'])
self.assertEqual(1, metadatas[0]['offset'])
self.assertTrue(
timezone.parse(metadatas[0]['last_log_timestamp']) > ts)
def test_read_timeout(self):
ts = pendulum.now().subtract(minutes=5)
self.es.delete(index=self.index_name, doc_type=self.doc_type, id=1)
logs, metadatas = self.es_task_handler.read(
self.ti, 1, {
'offset': 0,
'last_log_timestamp': str(ts),
'end_of_log': False
})
self.assertEqual(1, len(logs))
self.assertEqual(len(logs), len(metadatas))
self.assertEqual([''], logs)
self.assertTrue(metadatas[0]['end_of_log'])
# offset should be initialized to 0 if not provided.
self.assertEqual(0, metadatas[0]['offset'])
self.assertTrue(
timezone.parse(metadatas[0]['last_log_timestamp']) == ts)
def test_local_run(self):
args = create_mock_args(
task_id='print_the_context',
dag_id='example_python_operator',
subdir='/root/dags/example_python_operator.py',
interactive=True,
execution_date=timezone.parse('2018-04-27T08:39:51.298439+00:00'))
reset(args.dag_id)
with patch('argparse.Namespace', args) as mock_args:
run(mock_args)
dag = get_dag(mock_args)
task = dag.get_task(task_id=args.task_id)
ti = TaskInstance(task, args.execution_date)
ti.refresh_from_db()
state = ti.current_state()
self.assertEqual(state, State.SUCCESS)
def test_read_nonexistent_log(self):
ts = pendulum.now()
# In ElasticMock, search is going to return all documents with matching index
# and doc_type regardless of match filters, so we delete the log entry instead
# of making a new TaskInstance to query.
self.es.delete(index=self.index_name, doc_type=self.doc_type, id=1)
logs, metadatas = self.es_task_handler.read(
self.ti, 1, {
'offset': 0,
'last_log_timestamp': str(ts),
'end_of_log': False
})
self.assertEqual(1, len(logs))
self.assertEqual(len(logs), len(metadatas))
self.assertEqual([''], logs)
self.assertFalse(metadatas[0]['end_of_log'])
self.assertEqual(0, metadatas[0]['offset'])
# last_log_timestamp won't change if no log lines read.
self.assertTrue(
timezone.parse(metadatas[0]['last_log_timestamp']) == ts)
def test_test(self):
"""Test the `airflow test` command"""
args = create_mock_args(task_id='print_the_context',
dag_id='example_python_operator',
subdir=None,
execution_date=timezone.parse('2018-01-01'))
saved_stdout = sys.stdout
try:
sys.stdout = out = StringIO()
cli.test(args)
output = out.getvalue()
# Check that prints, and log messages, are shown
self.assertIn(
'Done. Returned value was: Whatever you return gets printed in the logs',
output)
self.assertIn(
"'example_python_operator__print_the_context__20180101'",
output)
finally:
sys.stdout = saved_stdout
def parse_execution_date(execution_date_str):
"""
Parse execution date string to datetime object.
"""
return timezone.parse(execution_date_str)