def test_dag_cached() -> None:
"""Test that DAG caching works."""
serv = MockServer()
with Fun(serv, defaults=options(distributed=False)):
dat = put(b"bla bla")
step1 = morph(lambda x: x.decode().upper().encode(), dat)
step2b = shell("echo 'not'", inp=dict(file1=step1))
merge = shell("cat file1 file2",
inp=dict(file1=step1, file2=step2b.stdout),
out=["file2"])
execute(merge)
with Fun(serv, defaults=options(distributed=False, evaluate=False)):
# Same as above, should run through with no evaluation
dat = put(b"bla bla")
step1 = morph(lambda x: x.decode().upper().encode(), dat)
step2b = shell("echo 'not'", inp=dict(file1=step1))
merge = shell("cat file1 file2",
inp=dict(file1=step1, file2=step2b.stdout),
out=["file2"])
execute(merge)
with Fun(serv, defaults=options(distributed=False, evaluate=False)):
dat = put(b"bla bla")
step1 = morph(lambda x: x.decode().upper().encode(), dat)
# DIFFERENT HERE: Trigger re-evaluation and raise
step2b = shell("echo 'knot'", inp=dict(file1=step1))
merge = shell("cat file1 file2",
inp=dict(file1=step1, file2=step2b.stdout),
out=["file2"])
with pytest.raises(RuntimeError):
execute(merge)
def test_toposort() -> None:
"""Test that we can topologically sort the subset."""
with Fun(MockServer(), options(distributed=False)) as db:
dat = put(b"bla bla")
step1 = morph(capitalize, dat)
step2 = shell("cat file1 file2", inp=dict(file1=step1, file2=dat))
# random not included ops
stepA = shell("echo 'bla'")
_ = concat(dat, dat)
_ = morph(capitalize, b"another word")
final = shell("cat file1 file2",
inp={
"file1": stepA.stdout,
"file2": step2.stdout
})
ops = _p._parametrize_subgraph(db, {"input": dat},
{"output": final.stdout})
edges = _p._subgraph_edges(db, ops)
sorted_ops = _p._subgraph_toposort(ops, edges)
assert sorted_ops[0] == step1.parent
assert sorted_ops[1] == step2.hash
assert sorted_ops[2] == final.hash
def test_error_propagation_morph() -> None:
"""Test propagation of errors."""
with Fun(MockServer()):
db, store = get_connection()
s1 = funsies.shell("cp file1 file3",
inp=dict(file1="bla"),
out=["file2"])
def fun_strict(inp: bytes) -> bytes:
return inp
def fun_lax(inp: Result[bytes]) -> bytes:
return b"bla bla"
s2 = funsies.morph(fun_strict, s1.out["file2"])
s3 = funsies.morph(fun_lax, s1.out["file2"])
s4 = funsies.morph(fun_lax, s1.out["file2"], strict=False)
run_op(db, store, s1.op.hash)
run_op(db, store, s2.parent)
out = funsies.take(s2, strict=False)
assert isinstance(out, Error)
assert out.source == s1.op.hash
print(s3.parent)
run_op(db, store, s3.parent)
out = funsies.take(s3, strict=False)
assert isinstance(out, Error)
assert out.source == s1.op.hash
run_op(db, store, s4.parent)
out = funsies.take(s4)
assert out == b"bla bla"
def test_dag_dump() -> None:
"""Test simple DAG dump to file."""
with Fun(MockServer(), options(distributed=False)) as db:
dat = put(b"bla bla")
dat2 = put(b"blaXbla")
errorstep = morph(raises, dat2)
step1 = morph(upper, dat)
step2 = shell("cat file1 file2", inp=dict(file1=step1, file2=dat))
step2b = utils.concat(step2.stdout, errorstep, strict=False)
step3 = shell("cat file1", inp=dict(file1=step2b))
step4 = shell("cat file1", inp=dict(file1=step1))
step4b = shell("cat file2", inp=dict(file2=step4.stdout))
out = utils.concat(step1, dat, step2.stdout, step3.stdout)
_dag.build_dag(db, out.hash)
execute(step2b)
execute(step4b)
wait_for(step4b, 1.0)
reset(step4)
nodes, artefacts, labels, links = _graphviz.export(
db, [out.hash, step4b.hash])
dot = _graphviz.format_dot(nodes, artefacts, labels, links,
[out.hash, step4b.hash])
assert len(dot) > 0
assert len(nodes) == 8
assert len(labels) == 8
# TODO pass through dot for testing?
with open("g.dot", "w") as f:
f.write(dot)
def test_subgraph() -> None:
"""Test that we can isolate the required operators for parametrization."""
with Fun(MockServer(), options(distributed=False)) as db:
dat = put(b"bla bla")
step1 = morph(capitalize, dat)
step2 = shell("cat file1 file2", inp=dict(file1=step1, file2=dat))
# random not included ops
stepA = shell("echo 'bla'")
_ = concat(dat, dat)
_ = morph(capitalize, b"another word")
final = shell("cat file1 file2",
inp={
"file1": stepA.stdout,
"file2": step2.stdout
})
ops = _p._parametrize_subgraph(db, {"input": dat},
{"output": final.stdout})
assert len(ops) == 3
assert step1.parent in ops
assert step2.hash in ops
assert final.hash in ops
# get edges
edges = _p._subgraph_edges(db, ops)
print(edges)
def test_rm() -> None:
"""Test rm."""
with Fun(MockServer(), options(distributed=False)):
dat = ui.put("bla bla")
# removing const artefact raises
with pytest.raises(AttributeError):
ui.reset(dat)
ui.take(dat)
def upper(x: str) -> str:
return x.upper()
m1 = fp.morph(upper, dat)
m2 = fp.morph(lambda x: x + x, m1)
ui.execute(m2)
assert ui.take(m2) == "BLA BLABLA BLA"
ui.reset(m1)
with pytest.raises(UnwrapError):
# deletion works
ui.take(m1)
with pytest.raises(UnwrapError):
# and it's recursive
ui.take(m2)
# re run
ui.execute(m2)
assert ui.take(m2) == "BLA BLABLA BLA"
def test_subdag() -> None:
"""Test that subdags execute properly."""
def cap(inp: bytes) -> bytes:
return inp.upper()
def map_reduce(
inputs: dict[str, bytes]) -> dict[str, _graph.Artefact[bytes]]:
"""Basic map reduce."""
inp_data = inputs["inp"].split(b" ")
out: list[_graph.Artefact[bytes]] = []
for el in inp_data:
out += [morph(cap, el, opt=options(distributed=False))]
return {"out": concat(*out, join="-")}
with Fun(MockServer(), defaults=options(distributed=False)) as db:
dat = put(b"bla bla lol what")
inp = {"inp": dat}
cmd = _subdag.subdag_funsie(map_reduce, {"inp": Encoding.blob},
{"out": Encoding.blob})
operation = _graph.make_op(db, cmd, inp, options())
out = _graph.Artefact[bytes].grab(db, operation.out["out"])
final = shell(
"cat file1 file2",
inp=dict(file1=out, file2=b"something"),
)
execute(final)
data = take(final.stdout)
assert data == b"BLA-BLA-LOL-WHATsomething"
def test_template() -> None:
"""Basic test of chevron templating."""
with Fun(MockServer()):
db, store = get_connection()
t = "Hello, {{ mustache }}!"
result = template(t, {"mustache": "world"})
run_op(db, store, result.parent)
assert take(result) == b"Hello, world!"
def test_store_cache() -> None:
"""Test store for caching."""
with Fun(MockServer()):
s = ui.put("bla bla")
s2 = ui.put(b"bla bla")
assert s != s2 # type:ignore
assert ui.take(s) == "bla bla"
assert ui.take(s2) == b"bla bla"
def test_store_takeout() -> None:
"""Test store for caching."""
with Fun(MockServer()):
s = ui.put(3)
with tempfile.NamedTemporaryFile() as f:
ui.takeout(s, f.name)
with open(f.name, "r") as f2:
assert json.loads(f2.read()) == 3
def test_truncate() -> None:
"""Test truncation."""
with Fun(MockServer()):
db, store = get_connection()
inp = "\n".join([f"{k}" for k in range(10)])
dat1 = put(inp.encode())
trunc = utils.truncate(dat1, 2, 3)
run_op(db, store, trunc.parent)
assert take(trunc) == ("\n".join(inp.split("\n")[2:-3])).encode()
def test_store_cache_jsons() -> None:
"""Test store for caching."""
with Fun(MockServer()):
# more complex data types
li = ui.put([1, 2, 3])
assert ui.take(li) == [1, 2, 3]
di = ui.put({"fun": 3})
assert ui.take(di) == {"fun": 3}
def test_wait() -> None:
"""Test waiting on things."""
with Fun(MockServer()):
db, store = _context.get_connection()
s = ui.shell("cp file1 file2", inp={"file1": "wawa"}, out=["file2"])
with pytest.raises(TimeoutError):
ui.wait_for(s.stdout, timeout=0)
run_op(db, store, s.op.hash)
ui.wait_for(s.stdout, timeout=0)
ui.wait_for(s, timeout=0)
def test_not_generated() -> None:
"""What happens when an artefact is not generated?"""
with Fun(MockServer()):
db, store = get_connection()
s = funsies.shell("cp file1 file2",
inp=dict(file1="bla"),
out=["file3"])
run_op(db, store, s.op.hash)
assert funsies.take(s.returncode) == 0
with pytest.raises(UnwrapError):
funsies.take(s.out["file3"])
def test_shell_run() -> None:
"""Test shell command."""
with Fun(MockServer()):
db, store = _context.get_connection()
s = ui.shell("cp file1 file2", inp={"file1": b"wawa"}, out=["file2"])
run_op(db, store, s.hash)
assert _graph.get_data(db, store, s.stderr) == b""
assert _graph.get_data(db, store, s.returncode) == 0
assert _graph.get_data(db, store, s.inp["file1"]) == b"wawa"
assert _graph.get_data(db, store, s.stdout) == b""
assert ui.take(s.out["file2"]) == b"wawa"
def test_parametric_store_recall() -> None:
"""Test storing and recalling parametrics."""
serv = MockServer()
with Fun(serv, options(distributed=False)):
a = put(3)
b = put(4)
s = reduce(lambda x, y: x + y, a, b)
s2 = morph(lambda x: 3 * x, s)
execute(s2)
assert take(s2) == 21
# parametrize
p.commit("math", inp=dict(a=a, b=b), out=dict(s=s, s2=s2))
with Fun(serv, options(distributed=False)):
out = p.recall("math", inp=dict(a=5, b=8))
execute(out["s2"])
assert take(out["s2"]) == 39
def test_dag_execute() -> None:
"""Test execution of a _dag."""
with Fun(MockServer(), defaults=options(distributed=False)):
dat = put(b"bla bla")
step1 = morph(lambda x: x.decode().upper().encode(), dat)
step2 = shell("cat file1 file2", inp=dict(file1=step1, file2=dat))
output = step2.stdout
# make queue
execute(output)
out = take(output)
assert out == b"BLA BLAbla bla"
def test_morph() -> None:
"""Test store for caching."""
with Fun(MockServer()):
db, store = _context.get_connection()
dat = ui.put(b"bla bla")
morph = fp.morph(lambda x: x.decode().upper().encode(), dat)
run_op(db, store, morph.parent)
assert ui.take(morph) == b"BLA BLA"
dat = ui.put("bla bla")
morph = fp.morph(lambda x: x.upper(), dat, name="CAPITALIZE_THIS")
run_op(db, store, morph.parent)
assert ui.take(morph) == "BLA BLA"
def test_shell_norun() -> None:
"""Test run on a shell command that didn't run."""
with Fun(MockServer()):
cmd = shell("cat file1", inp={"file1": b"bla bla"}, out=["bla"])
with tempfile.TemporaryDirectory() as d:
debug.shell(cmd, d)
n = os.listdir(d)
assert "input_files" in n
assert "output_files" in n
with open(os.path.join(d, "errors.json"), "r") as f:
assert "NotFound" in f.read()
def test_concat() -> None:
"""Test concatenation."""
with Fun(MockServer()):
db, store = get_connection()
dat1 = put(b"bla")
dat2 = put(b"bla")
cat = utils.concat(dat1, dat2)
run_op(db, store, cat.parent)
assert take(cat) == b"blabla"
cat = utils.concat(dat1, dat1, dat1, join=b" ")
run_op(db, store, cat.parent)
assert take(cat) == b"bla bla bla"
def test_parametric_store_recall_optional() -> None:
"""Test storing a parametric with optional parameters."""
serv = MockServer()
with Fun(serv, options(distributed=False)):
a = put(3)
b = put("fun")
s = reduce(lambda x, y: x * y, a, b)
s2 = morph(lambda x: x.upper(), s)
# parametrize
p.commit("fun", inp=dict(a=a, b=b), out=dict(s=s2))
with Fun(serv, options(distributed=False)):
out = p.recall("fun", inp=dict(a=5))
execute(out["s"])
assert take(out["s"]) == "FUNFUNFUNFUNFUN"
# nested
out = p.recall("fun", inp=dict(b="lol"))
out = p.recall("fun", inp=dict(b=out["s"], a=2))
execute(out["s"])
assert take(out["s"]) == "LOLLOLLOLLOLLOLLOL"
def test_reduce() -> None:
"""Test store for caching."""
with Fun(MockServer()):
db, store = _context.get_connection()
dat = ui.put("bla bla")
morph = fp.morph(lambda x: x.upper(), dat)
def join(x: str, y: str) -> str:
return x + y
red = fp.reduce(join, morph, dat)
run_op(db, store, morph.parent)
run_op(db, store, red.parent)
assert ui.take(red) == "BLA BLAbla bla"
def test_shell_run2() -> None:
"""Test shell command output side cases."""
with Fun(MockServer()):
db, store = _context.get_connection()
s = ui.shell("cp file1 file2", "cat file2", inp={"file1": b"wawa"})
run_op(db, store, s.hash)
assert _graph.get_data(db, store, s.inp["file1"]) == b"wawa"
with pytest.raises(Exception):
_graph.get_data(db, store, s.stdout)
with pytest.raises(Exception):
_graph.get_data(db, store, s.stderr)
with pytest.raises(Exception):
_graph.get_data(db, store, s.returncode)
assert ui.take(s.stdouts[1]) == b"wawa"
def test_dag_execute_same_root() -> None:
"""Test execution of two dags that share the same origin."""
with Fun(MockServer(), defaults=options(distributed=False)):
dat = put(b"bla bla")
step1 = morph(lambda x: x.decode().upper().encode(), dat)
step2 = shell("cat file1 file2", inp=dict(file1=step1, file2=dat))
step2b = shell("cat file1", inp=dict(file1=step1))
execute(step2)
out = take(step2.stdout)
assert out == b"BLA BLAbla bla"
execute(step2b)
out = take(step2b.stdout)
assert out == b"BLA BLA"
def test_error_propagation() -> None:
"""Test propagation of errors."""
with Fun(MockServer()):
db, store = get_connection()
s1 = funsies.shell("cp file1 file3",
inp=dict(file1="bla"),
out=["file2"])
s2 = funsies.shell("cat file1 file2",
inp=dict(file1="a file", file2=s1.out["file2"]))
run_op(db, store, s1.op.hash)
run_op(db, store, s2.op.hash)
out = funsies.take(s2.stdout, strict=False)
print(out)
assert isinstance(out, Error)
assert out.source == s1.op.hash
def test_parametric() -> None:
"""Test that parametric DAGs work."""
with Fun(MockServer(), options(distributed=False)) as db:
dat = put(b"bla bla")
step1 = morph(capitalize, dat)
step2 = shell("cat file1 file2", inp=dict(file1=step1, file2=dat))
final = shell("cat file1 file3",
inp={
"file1": step1,
"file3": step2.stdout
})
param = _p.make_parametric(db, "param", {"input": dat},
{"output": final.stdout})
param2 = _p.Parametric.grab(db, param.hash)
assert param == param2
def test_dag_large() -> None:
"""Test that DAG building doesn't do extra work for large operations."""
with Fun(MockServer()) as db:
outputs = []
for i in range(100):
dat = put(f"bla{i}".encode())
step1 = morph(lambda x: x.decode().upper().encode(), dat)
step2 = shell(
"cat file1 file2",
inp=dict(file1=step1, file2="something"),
out=["file2"],
)
outputs += [concat(step1, step1, step2.stdout, join=b" ")]
final = concat(*outputs, join=b"\n")
_dag.build_dag(db, final.hash)
assert len(_dag._dag_dependents(db, final.hash, hash_t("root"))) == 100
def test_shell_run() -> None:
"""Test run on a shell command."""
with Fun(MockServer()):
db, store = get_connection()
cmd = shell("cat file1", inp={"file1": b"bla bla"}, out=["bla"])
_ = run_op(db, store, cmd.hash)
with tempfile.TemporaryDirectory() as d:
debug.shell(cmd, d)
n = os.listdir(d)
assert "stdout0" in n
assert "stderr0" in n
assert "input_files" in n
assert "output_files" in n
with open(os.path.join(d, "errors.json"), "r") as f:
assert "MissingOutput" in f.read()
def test_template_complicated() -> None:
"""Test templating with funky types."""
with Fun(MockServer()):
db, store = get_connection()
t = "wazzaa, {{ mustache }}!"
result = template(t, {"mustache": put(b"people")})
run_op(db, store, result.parent)
assert take(result) == b"wazzaa, people!"
t = "{{a}}{{b}}{{c}}"
result = template(t, dict(a=2, b="cool", c="4me"))
run_op(db, store, result.parent)
assert take(result) == b"2cool4me"
t = ""
result = template(t, dict(a=2, b="cool", c="4me"))
run_op(db, store, result.parent)
assert take(result) == b""
def test_artefact() -> None:
"""Test artefact debug."""
with Fun(MockServer()) as db:
cmd = shell("cat file1", inp={"file1": b"bla bla"}, out=["bla"])
with tempfile.TemporaryDirectory() as d:
debug.artefact(cmd.stdout, d, connection=db)
n = os.listdir(d)
assert "metadata.json" in n
assert "error.json" in n
assert "data" not in n
with tempfile.TemporaryDirectory() as d:
debug.artefact(cmd.inp["file1"], d, connection=db)
n = os.listdir(d)
assert "metadata.json" in n
assert "error.json" not in n
assert "data" in n