def test_word_count_inline(self):
with dbnd_config(disable_tracker_api()):
assert_run_task(
word_count_inline.t(
text=TEXT_FILE,
task_version=str(random.random()),
override=conf_override,
))
def test_word_count_pyspark(self):
with dbnd_config(disable_tracker_api()):
actual = WordCountPySparkTask(
text=TEXT_FILE,
task_version=str(random.random()),
override=conf_override,
)
actual.dbnd_run()
print(target(actual.counters.path, "part-00000").read())
def test_word_spark_with_error(self):
with dbnd_config(disable_tracker_api()):
actual = WordCountThatFails(
text=TEXT_FILE,
task_version=str(random.random()),
override=conf_override,
)
with pytest.raises(DatabandRunError):
actual.dbnd_run()
def __init__(self, dag):
self.dag = dag
self.dbnd_airflow_name = {}
config_store = self.get_and_process_dbnd_dag_config()
with dbnd_config(
config_values=config_store, source="airflow"
) as current_config:
self.dbnd_context = DatabandContext(name="airflow__%s" % self.dag.dag_id)
with DatabandContext.context(_context=self.dbnd_context):
# we need databand context to update config first
self.dbnd_config_layer = current_config.config_layer
def test_word_count_with_hook_raise(self):
_config = conf_override.copy()
# add post submit hook location to the livy config
_config["livy"][
"job_submitted_hook"] = "test_livy_spark_task.hook_with_raise"
with dbnd_config(disable_tracker_api()):
with pytest.raises(DatabandRunError):
WordCountTask(text=TEXT_FILE,
task_version=str(random.random()),
override=_config).dbnd_run()
def test_word_count_with_job_submitted_hook_side_effect(self):
_config = conf_override.copy()
global STATE
STATE = False
# add job submitted hook location to the livy config
_config["livy"][
"job_submitted_hook"] = "test_livy_spark_task.side_effect"
with dbnd_config(disable_tracker_api()):
WordCountTask(text=TEXT_FILE,
task_version=str(random.random()),
override=_config).dbnd_run()
# calling the hook will call side_effect which will change the STATE to True
assert STATE
def test_word_count_with_job_status_hook_side_effect_3_times(self):
_config = conf_override.copy()
global CALL_COUNT
CALL_COUNT = 0
# add job status hook location to the livy config
_config["livy"][
"job_status_hook"] = "test_livy_spark_task.side_effect_status_update"
with dbnd_config(disable_tracker_api()):
WordCountTask(text=TEXT_FILE,
task_version=str(random.random()),
override=_config).dbnd_run()
# calling the hook will call side_effect which will change the STATE to True
assert (
CALL_COUNT >= 3
), "Status update hook should be called at least 3 times (queued -> pending -> success)"
def test_word_count_with_hook(self, mock_channel_tracker):
_config = conf_override.copy()
# add post submit hook location to the livy config
_config["livy"][
"job_submitted_hook"] = "test_livy_spark_task.add_proxy_as_external_link"
with dbnd_config(disable_tracker_api()):
WordCountTask(text=TEXT_FILE,
task_version=str(random.random()),
override=_config).dbnd_run()
calls = [
call for call in mock_channel_tracker.call_args_list
if call.args[0].__name__ == "save_external_links"
]
assert len(calls) == 1
assert calls[0].kwargs["external_links_dict"] == {
"proxy":
"https://proxy.proxy.local:12345/gateway//resource?scheme=http&host=livy&port=8998"
}
