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]