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" }