def test_if(self):
import tributary.streaming as ts
assert ts.run(
ts.Print(
ts.If(ts.Func(conditionals), ts.Func(if_stream),
ts.Func(else_stream)))) == [1, 2, 3]
def test_reduce_stateful(self):
def func1():
yield 0
yield 2
yield 4
def func2():
yield 1
yield 3
yield 5
def func3():
yield 2
yield 4
yield 6
def reduce(node1_value, node2_value, node3_value, reducer_node):
if not hasattr(reducer_node, "state"):
# on first call, make sure node tracks state
reducer_node.set("state", {"n1": None, "n2": None, "n3": None})
if node1_value is not None:
reducer_node.state["n1"] = node1_value
if node2_value is not None:
reducer_node.state["n2"] = node2_value
if node3_value is not None:
reducer_node.state["n3"] = node3_value
return reducer_node.state
n1 = ts.Func(func1)
n2 = ts.Func(func2)
n3 = ts.Func(func3)
r = ts.Reduce(n1, n2, n3, reducer=reduce, inject_node=True)
out = ts.run(r)
print(out)
assert out == [
{
"n1": 0,
"n2": 1,
"n3": 2
},
{
"n1": 2,
"n2": 3,
"n3": 4
},
{
"n1": 4,
"n2": 5,
"n3": 6
},
]
def test_merge(self):
def func1():
yield 1
yield 3
def func2():
yield 2
yield 4
yield 6
out = ts.Merge(ts.Print(ts.Func(func1)), ts.Print(ts.Func(func2)))
assert ts.run(out) == [(1, 2), (3, 4)]
def test_fixed_map(self):
def func():
yield [1, 2, 3]
yield [4, 5, 6]
out = ts.Reduce(*[x + 1 for x in ts.Func(func).map(3)])
assert ts.run(out) == [(2, 3, 4), (5, 6, 7)]
def test_diff(self):
ret = ts.run(ts.Diff(ts.Func(func2)))
vals = [1, 2, 0, 5, 4]
comp = [None] + [vals[i] - vals[i - 1] for i in range(1, 5)]
assert ret[0] is None
for i, x in enumerate(ret[1:]):
assert (x - comp[i + 1]) < 0.001
def test_reduce(self):
def func1():
yield 1
yield 4
def func2():
yield 2
yield 5
yield 7
def func3():
yield 3
yield 6
yield 8
out = ts.Reduce(ts.Func(func1), ts.Func(func2), ts.Func(func3))
assert ts.run(out) == [(1, 2, 3), (4, 5, 6)]
def test_http(self):
"""Test http server"""
def func():
yield "x"
yield "y"
yield "z"
out = ts.HTTPSink(ts.Func(func), url="http://localhost:8080")
assert len(ts.run(out)) == 3
def test_socketio(self):
"""Test socketio streaming"""
def func():
yield "a"
yield "b"
yield "c"
out = ts.SocketIOSink(ts.Func(func), url="http://localhost:8069")
assert ts.run(out) == ["a", "b", "c"]
def test_interval(self):
def reducer(short, long, node):
if not node.has("state"):
node.set("state", {"long": 0, "short": 0})
if long:
node.state["long"] += long
if short:
node.state["short"] += short
return node.state.copy()
async def short():
await sleep(1)
return 1
async def long():
await sleep(2)
return 2
def interval(func, time=1):
task = None
async def _ret():
nonlocal task
if task is None:
task = asyncio.ensure_future(func())
if not task.done():
await sleep(time)
return None
result = task.result()
task = None
return result
return _ret
short_node = ts.Func(interval(short, 1), count=5)
long_node = ts.Func(interval(long, 1), count=5)
out = ts.Reduce(short_node,
long_node,
reducer=reducer,
inject_node=True).print()
ts.run(out)
def test_websocket_server(self):
"""Test websocket server"""
def func():
yield "x"
yield "y"
yield "z"
out = ts.WebSocketServerSink(ts.Func(func), port=1234)
assert len(ts.run(out)) == 3
def test_func_generator(self):
def func():
yield 1
yield 2
yield 3
yield 4
yield 5
t = ts.Func(func)
assert ts.run(t) == [1, 2, 3, 4, 5]
def test_kafka(self):
"""Test streaming with Kafka"""
def func():
yield "a"
yield "b"
yield "c"
out = ts.KafkaSink(ts.Func(func), servers="localhost:9092", topic="tributary")
assert ts.run(out) == ["a", "b", "c"]
def test_websocket(self):
"""Test websocket streaming"""
def func():
yield "x"
yield "y"
yield "z"
out = ts.WebSocketSink(ts.Func(func), url="ws://localhost:8080", response=True)
assert len(ts.run(out)) == 3
def test_function_order(self):
def func(*args):
for _ in range(5):
yield _
lazy_node = tl.Node(value=func) + 5
# 5 6 7 8 9
streaming_node = ts.Print(ts.Func(func) - 5, "streaming:")
# -5 -4 -3 -2 -1
out = streaming_node + ts.Print(lazy_node, "lazy:")
x = ts.run(out)
# 0 2 4 6 8
print(x)
assert x == [0, 2, 4, 6, 8]
def test_value_order(self):
lazy_node = tl.Node(value=0)
# 5 6 7 8 9
def func(lazy_node=lazy_node):
for _ in range(5):
yield _
lazy_node.setValue(_ + 1)
streaming_node = ts.Print(ts.Func(func) - 5, "streaming:")
# -5 -4 -3 -2 -1
out = streaming_node + ts.Print(lazy_node + 5, "lazy:")
x = ts.run(out)
# 0 2 4 6 8
print(x)
assert x == [0, 2, 4, 6, 8]
def test_value(self):
def func(*args):
for _ in range(5):
lazy_node.setValue(_ + 1)
yield _
lazy_node = tl.Node(value=0)
# 5 6 7 8 9
streaming_node = ts.Print(ts.Func(func) - 5, "streaming:")
# -5 -4 -3 -2 -1
out = ts.Print(t.LazyToStreaming(lazy_node) + 5,
"lazy:") + streaming_node
x = ts.run(out)
# 0 2 4 6 8
print(x)
assert x == [0, 2, 4, 6, 8]
def test_pg(self):
def func():
yield 1
yield 2
yield 3
def parser(data):
return ["INSERT INTO test(col1) VALUES ({});".format(data)]
out = ts.PostgresSink(
ts.Func(func),
query_parser=parser,
user="******",
database="postgres",
password="******",
host="localhost:5432",
)
assert len(ts.run(out)) == 3
def test_file(self):
file = os.path.abspath(
os.path.join(os.path.dirname(__file__), "test_file_data.json")
)
if os.path.exists(file):
os.remove(file)
def func():
yield 1
yield 2
yield 3
yield 4
def read_file(file):
with open(file, "r") as fp:
data = fp.read()
return [int(x) for x in data]
# Test that output is equal to what is read (generalized)
out = ts.FileSink(ts.Func(func), filename=file, json=True)
assert ts.run(out) == read_file(file)
def test_construct_streaming(self):
# adapted from https://gist.github.com/raddy/bd0e977dc8437a4f8276
# spot, strike, vol, days till expiry, interest rate, call or put (1,-1)
spot, strike, vol, dte, rate, cp = sy.symbols(
"spot strike vol dte rate cp")
T = dte / 260.0
N = syNormal("N", 0.0, 1.0)
d1 = (sy.ln(spot / strike) + (0.5 * vol**2) * T) / (vol * sy.sqrt(T))
d2 = d1 - vol * sy.sqrt(T)
TimeValueExpr = sy.exp(-rate * T) * (cp * spot * cdf(N)(cp * d1) -
cp * strike * cdf(N)(cp * d2))
PriceClass = ts.construct_streaming(TimeValueExpr)
def strikes():
strike = 205
while strike <= 220:
yield strike
strike += 2.5
price = PriceClass(
spot=tss.Const(210.59),
# strike=tss.Print(tss.Const(205), text='strike'),
strike=tss.Func(strikes, interval=1),
vol=tss.Const(14.04),
dte=tss.Const(4),
rate=tss.Const(0.2175),
cp=tss.Const(-1),
)
ret = tss.run(tss.Print(price._starting_node))
time.sleep(2)
print(ret)
assert len(ret) == 7
def test_process(self):
def func():
yield {"a": 1, "b": 2}
yield {"a": 2, "b": 4}
yield {"a": 3, "b": 6}
yield {"a": 4, "b": 8}
def _json(val):
return JSON.dumps(val)
cmd = "{} {} --1".format(
sys.executable, os.path.join(os.path.dirname(__file__), "echo.py"))
print(cmd)
ret = ts.run(
ts.Subprocess(ts.Func(func).print("in:"), cmd,
json=True).print("out:"))
print(ret)
assert ret == [
{
"a": 1,
"b": 2
},
{
"a": 2,
"b": 4
},
{
"a": 3,
"b": 6
},
{
"a": 4,
"b": 8
},
]
def test_replace(self):
out = ts.Func(func, replace=True) + ts.Func(func2)
assert ts.run(out) == [2, 13, 105, 1006]
def test_drop(self):
out = ts.Func(func, drop=True) + ts.Func(func2)
assert ts.run(out) == [2, 12, 104, 1006]
def test_normal(self):
out = ts.Func(func) + ts.Func(func2)
assert ts.run(out) == [2, 12, 103, 1004]
def test_window_fixed_size(self):
assert ts.run(ts.Window(ts.Func(func), size=2)) == [[1], [1, 2]]
def test_window_fixed_size_full_only(self):
assert ts.run(ts.Window(ts.Func(func), size=2,
full_only=True)) == [[1, 2]]
def test_window_any_size(self):
assert ts.run(ts.Window(ts.Func(func))) == [[1], [1, 2]]
def test_repeat(self):
out = ts.Func(func) + ts.Func(func2, repeat=True)
assert ts.run(out) == [2, 3, 13, 14, 105, 1006]
def test_delay(self):
out = ts.Delay(ts.Func(func), delay=5)
now = time.time()
ret = ts.run(out)
assert time.time() - now > 5
assert ret == [1, 2]
def test_unroll(self):
assert ts.run(ts.Unroll(ts.Func(func2))) == [1, 2, 3, 4]
def test_apply(self):
def square(val):
return val**2
assert ts.run(ts.Apply(ts.Func(func), func=square)) == [1, 4]
