def test_run_bad2(): class TestWriterTmp(AbstractWriter): @staticmethod def necessary_config(node_config): return set([]) def run(self, data_object): raise Exception("Deliberate error") # return data_object, False NodeFactory().register("TestWriterTmp", TestWriterTmp) config = { "implementation_config": { "writer_config": { "mywriter": { "class": "TestWriterTmp", "destinations": [] } } } } configuration = Configuration(None, is_dict_config=True, dict_config=config) runner = DagRunner(configuration) with pytest.raises(Exception) as e: runner.run() assert "Deliberate error" in str(e)
def test_run2(): config = { "implementation_config": { "reader_config": { "csv_reader": { "class": "CsvReader", "filename": "test/minimal.csv", "destinations": [], } } } } configuration = Configuration(None, is_dict_config=True, dict_config=config) runner = DagRunner(configuration) with LogCapture() as l: runner.run(dry_run=True) l.check( ("root", "INFO", "Taking nodes to run from default"), ( "root", "INFO", "DRY RUN 0: would run node csv_reader of type reader_config and class CsvReader", ), ("root", "INFO", "All done. Bye bye!"), )
def test_run_bad(): class TestWriterTmp(AbstractWriter): @staticmethod def necessary_config(node_config): return set([]) def run(self, data_object): return data_object, False NodeFactory().register("TestWriterTmp", TestWriterTmp) config = { "implementation_config": { "writer_config": { "mywriter": { "class": "TestWriterTmp", "destinations": [] } } } } configuration = Configuration(None, is_dict_config=True, dict_config=config) runner = DagRunner(configuration) # unregister this class del NodeFactory().name_dict["TestWriterTmp"] with pytest.raises(Exception) as e: runner.run() assert "Issue instantiating mywriter and class TestWriterTmp" in str(e)
def test_run_notification_error(): config = { "metadata": { "section_registry": ["phase1"], "notify_on_error": { "client": "SlackClient", "channel": "some-channel", "token": "slack-api-token", "member_id": "optional-key", }, }, "implementation_config": { "phase1": { "csv_reader": { "class": "CsvReader", "filename": "bad/path.csv" } } }, } configuration = Configuration(None, is_dict_config=True, dict_config=config) runner = DagRunner(configuration) mock_client = mock.Mock() mock_client.post_message = mock.Mock() mock_get_notification_client = mock.Mock(return_value=mock_client) path = "primrose.notification_utils.get_notification_client" with mock.patch(path) as mock_get_notification_client: with pytest.raises(Exception) as error: runner.run() assert mock_get_notification_client.post_message.call_count == 1
def test_run4(): class TestWriter(AbstractFileWriter): def __init__(self, configuration, instance_name): pass @staticmethod def necessary_config(node_config): return set([]) def run(self, data_object): terminate = True return data_object, terminate NodeFactory().register("TestWriter", TestWriter) config = { "implementation_config": { "reader_config": { "csv_reader": { "class": "CsvReader", "filename": "test/minimal.csv", "destinations": ["csv_writer"], } }, "writer_config": { "csv_writer": { "class": "TestWriter" } }, } } configuration = Configuration(None, is_dict_config=True, dict_config=config) runner = DagRunner(configuration) with LogCapture() as l: runner.run(dry_run=False) l.check( ("root", "INFO", "Taking nodes to run from default"), ( "root", "INFO", "received node csv_reader of type reader_config and class CsvReader", ), ("root", "INFO", "Reading test/minimal.csv from CSV"), ( "root", "INFO", "received node csv_writer of type writer_config and class TestWriter", ), ("root", "INFO", "Terminating early due to signal from csv_writer"), ("root", "INFO", "All done. Bye bye!"), )
def test_run6(): config = { "metadata": { "section_registry": ["phase1", "cleanup_config"], "notify_on_error": { "client": "SlackClient", "channel": "some-channel", "token": "slack-api-token", "member_id": "optional-key", }, }, "implementation_config": { "phase1": { "csv_reader": { "class": "CsvReader", "filename": "test/minimal.csv", "destinations": ["notification"], } }, "cleanup_config": { "notification": { "class": "ClientNotification", "client": "SlackClient", "channel": "some-channel", "token": "slack-api-token", "member_id": "optional-key", "message": "Yay! Sucess", } }, }, } configuration = Configuration(None, is_dict_config=True, dict_config=config) runner = DagRunner(configuration) with LogCapture() as l: runner.run(dry_run=True) l.check( ("root", "INFO", "Taking nodes to run from section_registry"), ( "root", "INFO", "DRY RUN 0: would run node csv_reader of type phase1 and class CsvReader", ), ( "root", "INFO", "DRY RUN 1: would run node notification of type cleanup_config and class ClientNotification", ), ("root", "INFO", "All done. Bye bye!"), )
def test_run(): config = { "implementation_config": { "reader_config": { "csv_reader": { "class": "CsvReader", "filename": "test/minimal.csv", "destinations": [], } } } } configuration = Configuration(None, is_dict_config=True, dict_config=config) runner = DagRunner(configuration) runner.run()
def test_run5(): config = { "metadata": { "section_registry": ["phase1", "phase2"] }, "implementation_config": { "phase1": { "csv_reader": { "class": "CsvReader", "filename": "test/minimal.csv", "destinations": ["csv_writer"], } }, "phase2": { "csv_writer": { "class": "CsvWriter", "key": "test_data", "dir": "cache", "filename": "test/unittest_similar_recipes.csv", } }, }, } configuration = Configuration(None, is_dict_config=True, dict_config=config) runner = DagRunner(configuration) with LogCapture() as l: runner.run(dry_run=True) l.check( ("root", "INFO", "Taking nodes to run from section_registry"), ( "root", "INFO", "DRY RUN 0: would run node csv_reader of type phase1 and class CsvReader", ), ( "root", "INFO", "DRY RUN 1: would run node csv_writer of type phase2 and class CsvWriter", ), ("root", "INFO", "All done. Bye bye!"), )
def test_run_pruned(): config = { "implementation_config": { "reader_config": { "read_data": { "class": "CsvReader", "filename": "data/tennis.csv", "destinations": ["conditional_node"], }, "conditional_node": { "class": "SimpleSwitch", "path_to_travel": "left", "destinations": ["left", "right"], }, }, "writer_config": { "left": { "class": "LoggingSuccess", "msg": "left node!", "level": "INFO", }, "right": { "class": "LoggingSuccess", "msg": "right node!", "level": "INFO", "destinations": ["right2"], }, "right2": { "class": "LoggingSuccess", "msg": "right node2!", "level": "INFO", }, }, } } configuration = Configuration(None, is_dict_config=True, dict_config=config) runner = DagRunner(configuration) with LogCapture() as l: runner.run() l.check( ("root", "INFO", "Taking nodes to run from default"), ( "root", "INFO", "received node read_data of type reader_config and class CsvReader", ), ("root", "INFO", "Reading data/tennis.csv from CSV"), ( "root", "INFO", "received node conditional_node of type reader_config and class SimpleSwitch", ), ("root", "INFO", "Skipping pruned node right"), ("root", "INFO", "Skipping pruned node right2"), ( "root", "INFO", "received node left of type writer_config and class LoggingSuccess", ), ("root", "INFO", "left node!"), ("root", "INFO", "All done. Bye bye!"), )