def test_methodcaller():
with BufferingNodeExecutionContext(methodcaller("swapcase")) as context:
context.write_sync("aaa", "bBb", "CcC")
assert context.get_buffer() == list(map(ensure_tuple, ["AAA", "BbB", "cCc"]))
with BufferingNodeExecutionContext(methodcaller("zfill", 5)) as context:
context.write_sync("a", "bb", "ccc")
assert context.get_buffer() == list(map(ensure_tuple, ["0000a", "000bb", "00ccc"]))
def test_methodcaller():
with BufferingNodeExecutionContext(methodcaller('swapcase')) as context:
context.write_sync('aaa', 'bBb', 'CcC')
assert context.get_buffer() == list(
map(ensure_tuple, ['AAA', 'BbB', 'cCc']))
with BufferingNodeExecutionContext(methodcaller('zfill', 5)) as context:
context.write_sync('a', 'bb', 'ccc')
assert context.get_buffer() == list(
map(ensure_tuple, ['0000a', '000bb', '00ccc']))
def test_fixedwindow():
with BufferingNodeExecutionContext(bonobo.FixedWindow(2)) as context:
context.write_sync(*range(10))
assert context.get_buffer() == [(0, 1), (2, 3), (4, 5), (6, 7), (8, 9)]
with BufferingNodeExecutionContext(bonobo.FixedWindow(2)) as context:
context.write_sync(*range(9))
assert context.get_buffer() == [(0, 1), (2, 3), (4, 5), (6, 7), (8, None)]
with BufferingNodeExecutionContext(bonobo.FixedWindow(1)) as context:
context.write_sync(*range(3))
assert context.get_buffer() == [(0, ), (1, ), (2, )]
def test_inherit_bag_tuple():
with BufferingNodeExecutionContext(append) as context:
context.set_input_fields(["message"])
context.write_sync(*messages)
assert context.get_output_fields() == ("message", "0")
assert context.get_buffer() == list(map(lambda x: x + ("!", ), messages))
def test_not_modified():
input_messages = [('foo', 'bar'), ('foo', 'baz')]
with BufferingNodeExecutionContext(useless) as context:
context.write_sync(*input_messages)
assert context.get_buffer() == input_messages
def test_set_fields():
with BufferingNodeExecutionContext(bonobo.SetFields(["x", "y"])) as context:
context.write_sync((1, 2))
output = context.get_buffer()
assert len(output) == 1
assert output[0]._fields == ("x", "y")
assert output[0].x == 1
assert output[0].y == 2
def test_cast_after_returning_custom_type():
def transform():
yield MyCustomType(1, 2, 3)
yield MyCustomType(4, 5, 6)
with BufferingNodeExecutionContext(transform) as context:
context.write_sync(EMPTY)
result = context.get_buffer()
assert tuple == type(result[0])
assert tuple == type(result[1])
assert MyCustomType == type(result[0][0])
assert MyCustomType == type(result[1][0])
with BufferingNodeExecutionContext(MyCustomType.as_tuple) as context:
context.write_sync(*result)
result = context.get_buffer()
assert MyBag == type(result[0])
assert MyBag == type(result[1])
def test_not_modified():
input_messages = [("foo", "bar"), ("foo", "baz")]
with BufferingNodeExecutionContext(useless) as context:
context.write_sync(*input_messages)
result = context.get_buffer()
print(result)
assert result == input_messages
def test_read_pickled_list_from_file(tmpdir):
fs, filename, services = pickle_tester.get_services_for_reader(tmpdir)
with BufferingNodeExecutionContext(PickleReader(filename),
services=services) as context:
context.write_sync(EMPTY)
output = context.get_buffer()
assert context.get_output_fields() == ("a", "b", "c")
assert output == [("a foo", "b foo", "c foo"), ("a bar", "b bar", "c bar")]
def test_read_pickled_list_from_file(tmpdir):
fs, filename, services = pickle_tester.get_services_for_reader(tmpdir)
with BufferingNodeExecutionContext(PickleReader(filename),
services=services) as context:
context.write_sync(EMPTY)
output = context.get_buffer()
assert context.get_output_fields() == ('a', 'b', 'c')
assert output == [('a foo', 'b foo', 'c foo'), ('a bar', 'b bar', 'c bar')]
def test_node_dict():
def f():
return {"id": 1, "name": "foo"}
with BufferingNodeExecutionContext(f) as context:
context.write_sync(EMPTY)
output = context.get_buffer()
assert len(output) == 1
assert output[0] == ({"id": 1, "name": "foo"},)
def g():
yield {"id": 1, "name": "foo"}
yield {"id": 2, "name": "bar"}
with BufferingNodeExecutionContext(g) as context:
context.write_sync(EMPTY)
output = context.get_buffer()
assert len(output) == 2
assert output[0] == ({"id": 1, "name": "foo"},)
assert output[1] == ({"id": 2, "name": "bar"},)
def test_file_reader(tmpdir):
fs, filename, services = txt_tester.get_services_for_reader(tmpdir)
with BufferingNodeExecutionContext(FileReader(path=filename),
services=services) as context:
context.write_sync(EMPTY)
output = context.get_buffer()
assert len(output) == 2
assert output[0] == ('Hello', )
assert output[1] == ('World', )
def test_node_dict():
def f():
return {'id': 1, 'name': 'foo'}
with BufferingNodeExecutionContext(f) as context:
context.write_sync(EMPTY)
output = context.get_buffer()
assert len(output) == 1
assert output[0] == ({'id': 1, 'name': 'foo'},)
def g():
yield {'id': 1, 'name': 'foo'}
yield {'id': 2, 'name': 'bar'}
with BufferingNodeExecutionContext(g) as context:
context.write_sync(EMPTY)
output = context.get_buffer()
assert len(output) == 2
assert output[0] == ({'id': 1, 'name': 'foo'},)
assert output[1] == ({'id': 2, 'name': 'bar'},)
def test_node_bytes():
def f():
return b'foo'
with BufferingNodeExecutionContext(f) as context:
context.write_sync(EMPTY)
output = context.get_buffer()
assert len(output) == 1
assert output[0] == (b'foo',)
def g():
yield b'foo'
yield b'bar'
with BufferingNodeExecutionContext(g) as context:
context.write_sync(EMPTY)
output = context.get_buffer()
assert len(output) == 2
assert output[0] == (b'foo',)
assert output[1] == (b'bar',)
def test_casts_after_output(factory: Callable, expected, expected_item0):
def transform():
yield factory()
yield factory()
with BufferingNodeExecutionContext(transform) as context:
context.write_sync(EMPTY)
result = context.get_buffer()
assert expected == type(result[0])
assert expected_item0 == type(result[0][0])
assert expected == type(result[1])
assert expected_item0 == type(result[1][0])
def test_node_tuple_dict():
def f():
return 'foo', 'bar', {'id': 1}
with BufferingNodeExecutionContext(f) as context:
context.write_sync(EMPTY)
output = context.get_buffer()
assert len(output) == 1
assert output[0] == ('foo', 'bar', {'id': 1})
def g():
yield 'foo', 'bar', {'id': 1}
yield 'foo', 'baz', {'id': 2}
with BufferingNodeExecutionContext(g) as context:
context.write_sync(EMPTY)
output = context.get_buffer()
assert len(output) == 2
assert output[0] == ('foo', 'bar', {'id': 1})
assert output[1] == ('foo', 'baz', {'id': 2})
def test_node_string():
def f():
return "foo"
with BufferingNodeExecutionContext(f) as context:
context.write_sync(EMPTY)
output = context.get_buffer()
assert len(output) == 1
assert output[0] == ("foo",)
def g():
yield "foo"
yield "bar"
with BufferingNodeExecutionContext(g) as context:
context.write_sync(EMPTY)
output = context.get_buffer()
assert len(output) == 2
assert output[0] == ("foo",)
assert output[1] == ("bar",)
def test_node_tuple_dict():
def f():
return "foo", "bar", {"id": 1}
with BufferingNodeExecutionContext(f) as context:
context.write_sync(EMPTY)
output = context.get_buffer()
assert len(output) == 1
assert output[0] == ("foo", "bar", {"id": 1})
def g():
yield "foo", "bar", {"id": 1}
yield "foo", "baz", {"id": 2}
with BufferingNodeExecutionContext(g) as context:
context.write_sync(EMPTY)
output = context.get_buffer()
assert len(output) == 2
assert output[0] == ("foo", "bar", {"id": 1})
assert output[1] == ("foo", "baz", {"id": 2})
def test_read_csv_from_file_kwargs(tmpdir):
fs, filename, services = csv_tester.get_services_for_reader(tmpdir)
with BufferingNodeExecutionContext(CsvReader(filename, **defaults),
services=services) as context:
context.write_sync(EMPTY)
assert context.get_buffer_args_as_dicts() == [{
'a': 'a foo',
'b': 'b foo',
'c': 'c foo',
}, {
'a': 'a bar',
'b': 'b bar',
'c': 'c bar',
}]
def test_read_csv_from_file_kwargs(tmpdir):
fs, filename, services = csv_tester.get_services_for_reader(tmpdir)
with BufferingNodeExecutionContext(CsvReader(filename, **defaults),
services=services) as context:
context.write_sync(EMPTY)
assert context.get_buffer_args_as_dicts() == [
{
"a": "a foo",
"b": "b foo",
"c": "c foo"
},
{
"a": "a bar",
"b": "b bar",
"c": "c bar"
},
]
def test_map_fields(input_, key, expected):
with BufferingNodeExecutionContext(bonobo.MapFields(lambda x: x ** 2, key)) as context:
context.write_sync(input_)
assert context.status == "-"
[got] = context.get_buffer()
assert expected == got
def test_map_fields_error():
with BufferingNodeExecutionContext(bonobo.MapFields(lambda x: x ** 2, lambda x: x == "c")) as context:
context.write_sync(tuple())
assert context.status == "!"
assert context.defunct
def test_inherit():
with BufferingNodeExecutionContext(append) as context:
context.write_sync(*messages)
assert context.get_buffer() == list(map(lambda x: x + ("!", ), messages))
def execute_node(self, node, *rows):
with BufferingNodeExecutionContext(node) as context:
context.write_sync(*rows)
return context.get_buffer()
