def test_get_stream_new_without_transport(self):
node_a, node_b = MockNode(), MockNode()
session_mock = MockSession(MockNetwork())
session_mock._id = "1"
session_mock.closed = False
session_mock.closing_mode = None
transport_mock = ZeroTransport(MockConnector())
def transport_factory(_):
return transport_mock
transport_factory_mock = mock.MagicMock(side_effect=transport_factory)
ZeroPipe.set_transport_factory(transport_factory_mock)
pipe = ZeroPipe(node_a, node_b)
ZeroStream.load = load
ZeroStream.save = mock.MagicMock()
stream = pipe.get_stream(session_mock)
ZeroStream.save.assert_called_once_with(save_dependency=False)
ZeroPipe._transport_factory.assert_called_once_with(pipe)
assert stream.transport == transport_mock
def test_push_many_return_piece_of_data(self):
node_a, node_b = MockNode(), MockNode()
session_mock = MockSession(MockNetwork())
session_mock._id = "1"
stream_mock = mock.MagicMock()
def get_stream(_):
return stream_mock
def lambda_dummy(x, *_, **__):
return x
pipe = ZeroPipe(node_a,
node_b,
filters=[FNFilter(lambda_dummy)],
filters_threshold=2)
pipe.get_stream = mock.MagicMock(side_effect=get_stream)
res = pipe.push(list(range(5)), session_mock, many=True, test_kw=1)
assert res is True
pipe.get_stream.assert_called_once_with(session_mock)
stream_mock.write.assert_called_once_with(list(range(2, 5)),
many=True,
test_kw=1)
def test_get_stream_exists(self):
node_a, node_b = MockNode(), MockNode()
session_mock = mock.MagicMock(spec=ZeroSession)
session_mock.id = "1"
session_mock.closed = False
session_mock.closing_mode = None
stream_mock = mock.MagicMock()
stream_cls_mock = mock.MagicMock(
side_effect=lambda *_, **__: stream_mock)
stream_cls_mock.load = mock.MagicMock(
side_effect=lambda *_, **__: stream_mock)
pipe = ZeroPipe(node_a, node_b)
@contextmanager
def patch():
old_load = ZeroStream.load
old_new = ZeroStream.__new__
try:
ZeroStream.load = lambda *_, **__: stream_mock
ZeroStream.__new__ = lambda *_, **__: stream_mock
yield ZeroStream
finally:
ZeroStream.load = old_load
ZeroStream.__new__ = old_new
with patch():
assert pipe.get_stream(session_mock) == stream_mock
def test_set_transport_factory_as_transport(self):
node_a, node_b = MockNode(), MockNode()
pipe = ZeroPipe(node_a, node_b)
transport = ZeroTransport(MockConnector())
pipe.set_transport_factory(transport)
assert pipe._transport_factory() == transport
def get_network():
hello_world_node = ZeroNode(
hello_world_op,
_id="hello_world_node_id", # set custom id (default: uuid4)
)
print_node = ZeroNode(print)
# --------> [print_node]
# |
# <in> ----> [hello_world_node] ----
# \
# --------> <out>
return ZeroNetwork([
ZeroPipe(ZeroNode.IN, hello_world_node),
ZeroPipe(hello_world_node, print_node),
ZeroPipe("hello_world_node_id", ZeroNode.OUT),
])
def get_network():
b = ZeroNode(RoundRobinBalancerOperator(), _id="balancer").save()
n1 = ZeroNode(PythonFunctionOperator(get_file_data), _id="node1").save()
n2 = ZeroNode(PythonFunctionOperator(get_file_data), _id="node2").save()
n3 = OwnZeroNode(PythonFunctionOperator(get_file_data), _id="node3").save()
bash_save = ZeroNode(BashPerPacketOperator(cmd_generator), _id="bash_save")
return ZeroNetwork([
ZeroPipe(ZeroNode.IN, b, _id="in-balancer"),
ZeroPipe("balancer", "/zero/nodes/node1", _id="balancer-node1"),
ZeroPipe("balancer", "/zero/nodes/node2", _id="balancer-node2"),
ZeroPipe("balancer", "/zero/nodes/own/node3", _id="balancer-node3"),
ZeroPipe("node1", bash_save, _id="node1-bash_save"),
ZeroPipe("node2", bash_save, _id="node2-bash_save"),
ZeroPipe("/zero/nodes/own/node3", bash_save, _id="node3-bash_save"),
ZeroPipe(n1, ZeroNode.OUT, _id="node1-out"),
ZeroPipe(n2, ZeroNode.OUT, _id="node2-out"),
ZeroPipe(n3, ZeroNode.OUT, _id="node3-out"),
])
def test_transport_factory_context(self):
node_a, node_b = MockNode(), MockNode()
pipe = ZeroPipe(node_a, node_b)
transport = mock.MagicMock()
prev_method = pipe._transport_factory
with pipe.transport_factory_context(transport):
assert pipe._transport_factory == transport
assert pipe._transport_factory == prev_method
def test_get_stream_new_when_sess_hard_closing(self):
node_a, node_b = MockNode(), MockNode()
session_mock = MockSession(MockNetwork())
session_mock._id = "1"
session_mock.closed = False
session_mock.closing_mode = ZeroSession.HARD_CLOSING_MODE
pipe = ZeroPipe(node_a, node_b)
with pytest.raises(TethysSessionClosed):
pipe.get_stream(session_mock)
def test_get_stream_new_when_sess_closed(self):
node_a, node_b = MockNode(), MockNode()
session_mock = MockSession(MockNetwork())
session_mock._id = "1"
session_mock.closed = True
session_mock.closing_mode = None
pipe = ZeroPipe(node_a, node_b)
with pytest.raises(TethysSessionClosed):
pipe.get_stream(session_mock)
def generate_network(name):
transport = ZeroTransport(transport_connectors_factory,
serializer=serializer,
deserializer=deserializer)
with ZeroPipe.transport_factory_context(transport):
req_http = ZeroNode(ReqOperator(req_http_op))
req_ftp = ZeroNode(ReqOperator(req_ftp_op))
save = ZeroNode(PythonFunctionOperator(save_op))
http_filter = RegexMatchFilter("^https?://.*")
ftp_filter = RegexMatchFilter("^ftp://.*")
return ZeroNetwork(
[
ZeroPipe(ZeroNode.IN, req_http, filters=[http_filter]),
ZeroPipe(ZeroNode.IN, req_ftp, filters=[ftp_filter]),
ZeroPipe(req_http, save),
ZeroPipe(req_ftp, save),
],
_id=name,
)
def test_set_transport_factory_and_create_pipe(self):
node_a, node_b = MockNode(), MockNode()
transport = ZeroTransport(MockConnector())
transport_factory = mock.MagicMock(
side_effect=lambda *_, **__: transport)
ZeroPipe.set_transport_factory(transport_factory)
pipe = ZeroPipe(node_a, node_b)
assert ZeroPipe._transport_factory == transport_factory
assert pipe.transport == transport
ZeroPipe._transport_factory.assert_called_once_with(pipe)
def get_network():
node1 = ZeroNode(PythonFunctionOperator(generate_data), _id="node1").save()
def f_odd(value):
return value % 2 == 1
def f_even(value):
return value % 2 == 0
odd_filter = SimpleJsonPathFilter("data.number", f_odd)
even_filter = SimpleJsonPathFilter("data.number", f_even)
custom_transport = ZeroTransport(
PersistQueueConnector(TMP_FOLDER, queue_engine="file"))
return ZeroNetwork([
ZeroPipe(
ZeroNode.IN,
node1,
_id="in-node1",
),
ZeroPipe(
"/zero/nodes/node1",
ZeroNode.OUT,
filters=[even_filter],
transport=custom_transport,
_id="even-out",
),
ZeroPipe(
node1,
ZeroNode.OUT,
filters=[odd_filter],
transport=custom_transport,
_id="odd-out",
),
])
def test_filter_data_packet_false(self):
node_a, node_b = MockNode(), MockNode()
session_mock = mock.MagicMock()
def f1(data_packet, *_, **__):
return data_packet
def f2(data_packet, *_, **__):
return data_packet / 2
pipe = ZeroPipe(node_a,
node_b,
filters=[FNFilter(f1), FNFilter(f2)],
filters_threshold=1.1)
assert pipe.filter_data_packet(2, session_mock) is False
def test_push_many(self):
node_a, node_b = MockNode(), MockNode()
session_mock = MockSession(MockNetwork())
session_mock._id = "1"
stream_mock = mock.MagicMock()
def get_stream(_):
return stream_mock
pipe = ZeroPipe(node_a, node_b)
pipe.get_stream = mock.MagicMock(side_effect=get_stream)
res = pipe.push([...], session_mock, many=True, test_kw=1)
assert res is True
pipe.get_stream.assert_called_once_with(session_mock)
stream_mock.write.assert_called_once_with([...], many=True, test_kw=1)
def test_get_stream_new_with_transport(self):
node_a, node_b = MockNode(), MockNode()
session_mock = MockSession(MockNetwork())
session_mock._id = "1"
session_mock.closed = False
session_mock.closing_mode = None
transport_mock = ZeroTransport(MockConnector())
pipe = ZeroPipe(node_a, node_b, transport=transport_mock)
ZeroStream.load = load
ZeroStream.save = mock.MagicMock()
stream = pipe.get_stream(session_mock)
ZeroStream.save.assert_called_once_with(save_dependency=False)
assert stream.transport == transport_mock
def test_pull(self):
node_a, node_b = MockNode(), MockNode()
session_mock = MockSession(MockNetwork())
session_mock._id = "1"
session_mock.closed = False
session_mock.closing_mode = ZeroSession.HARD_CLOSING_MODE
stream_mock = mock.MagicMock()
stream_mock.read = mock.MagicMock(
side_effect=lambda *_, **__: iter([("key", "value")]))
def get_stream(_):
return stream_mock
pipe = ZeroPipe(node_a, node_b)
pipe.get_stream = mock.MagicMock(side_effect=get_stream)
assert next(pipe.pull(session_mock, test_kw=1)) == "value"
pipe.get_stream.assert_called_once_with(session_mock)
stream_mock.read.assert_called_once_with(test_kw=1)
def test_push_many_return_empty(self):
node_a, node_b = MockNode(), MockNode()
session_mock = MockSession(MockNetwork())
session_mock._id = "1"
stream_mock = mock.MagicMock()
def get_stream(_):
return stream_mock
def lambda_null(*_, **__):
return 0
pipe = ZeroPipe(node_a, node_b, filters=[FNFilter(lambda_null)])
pipe.get_stream = mock.MagicMock(side_effect=get_stream)
res = pipe.push(list(range(5)), session_mock, many=True, test_kw=1)
assert res is False
pipe.get_stream.assert_not_called()
stream_mock.write.assert_not_called()
def test_filter_data_packet_empty(self):
node_a, node_b = MockNode(), MockNode()
session_mock = mock.MagicMock()
pipe = ZeroPipe(node_a, node_b, filters=[])
assert pipe.filter_data_packet(..., session_mock)