def test_complex_side_effects(self):
m, n = 5, 3
p = uberjob.Plan()
r = uberjob.Registry()
s = TestStore()
x = [p.call(lambda k: k * k, i) for i in range(m)]
w = [
p.call(
lambda k: s.write((s.read()
if s.get_modified_time() else 0) + k), x[i])
for i in range(m)
]
for a, b in (w[i:i + 2] for i in range(m - 1)):
p.add_dependency(a, b)
y = r.source(p, s)
for i in range(m):
p.add_dependency(w[i], y)
z = [p.call(lambda k: k + s.read(), i) for i in range(n)]
for i in range(n):
p.add_dependency(y, z[i])
for _ in range(2):
self.assertEqual(
uberjob.run(p, registry=r, output=z),
[a + sum(b * b for b in range(m)) for a in range(n)],
)
def build_plan(args):
plan = uberjob.Plan()
registry = uberjob.Registry()
f = registry.source(plan, JsonFileStore(os.path.join(args.input_directory, "input.json")))
c = plan.call(fahrenheit_to_celcius, f)
registry.add(c, JsonFileStore(os.path.join(args.output_directory, "output.json")))
return plan, registry
def build_plan(args):
plan = uberjob.Plan()
registry = uberjob.Registry()
f = registry.source(plan, JsonFileStore(args.input_path))
c = plan.call(fahrenheit_to_celcius, f)
registry.add(c, JsonFileStore(args.output_path))
return plan, registry
def test_fresh_time_advanced(self):
p = uberjob.Plan()
r = uberjob.Registry()
t0 = dt.datetime.now()
store_a = TestStore(7, modified_time=t0)
a = r.source(p, store_a)
store_c = TestStore()
b = p.call(store_c.write, a)
c = r.source(p, store_c)
p.add_dependency(b, c)
uberjob.run(p, registry=r)
self.assertEqual(store_a.read_count, 1)
self.assertEqual(store_c.read_count, 0)
self.assertEqual(store_c.write_count, 1)
self.assertGreater(store_c.modified_time, store_a.modified_time)
uberjob.run(p, registry=r, fresh_time=store_c.modified_time)
self.assertEqual(store_a.read_count, 1)
self.assertEqual(store_c.read_count, 0)
self.assertEqual(store_c.write_count, 1)
uberjob.run(p,
registry=r,
fresh_time=store_c.modified_time + dt.timedelta(seconds=1))
self.assertEqual(store_a.read_count, 2)
self.assertEqual(store_c.read_count, 0)
self.assertEqual(store_c.write_count, 2)
def test_registry_already_has_value_store(self):
plan = uberjob.Plan()
registry = uberjob.Registry()
x = plan.call(operator.add, 1, 2)
registry.add(x, TestStore())
with self.assertRaisesRegex(Exception,
r"The node already has a value store\."):
registry.add(x, TestStore())
def test_source_dependent_on_write(self):
p = uberjob.Plan()
r = uberjob.Registry()
x = p.call(lambda: 5)
r.add(x, TestStore())
y = r.source(p, r[x])
p.add_dependency(x, y)
self.assertEqual(uberjob.run(p, output=y, registry=r), 5)
def test_failed_to_read_from_empty_store_2(self):
p = uberjob.Plan()
r = uberjob.Registry()
x = r.source(p, TestStore())
y = r.source(p, TestStore())
p.add_dependency(x, y)
with self.assert_failed_to_read_from_empty_store():
uberjob.run(p, registry=r, output=y)
def test_registry_can_lie(self):
p = uberjob.Plan()
r = uberjob.Registry()
x = p.call(operator.add, 2, 3)
r.add(x, TestStore(7))
self.assertEqual(uberjob.run(p, output=x), 5)
self.assertEqual(r[x].read_count, 0)
self.assertEqual(r[x].write_count, 0)
self.assertEqual(uberjob.run(p, output=x, registry=r), 7)
def test_stack_frame_registry_add(self):
plan = uberjob.Plan()
registry = uberjob.Registry()
x = plan.call(operator.add, 2, 2)
stack_frame = get_stack_frame(1)
registry.add(x, TestStore(can_write=False))
with self.assert_call_exception(
expected_stack_frame=copy_with_line_offset(stack_frame, 1)):
uberjob.run(plan, registry=registry, output=x)
def test_has_a_cycle(self):
plan = uberjob.Plan()
registry = uberjob.Registry()
x = registry.source(plan, TestStore(1))
y = registry.source(plan, TestStore(2))
plan.add_dependency(x, y)
plan.add_dependency(y, x)
with self.assertRaises(nx.HasACycle):
uberjob.run(plan, registry=registry, output=y)
def test_render_dry_run():
plan = uberjob.Plan()
registry = uberjob.Registry()
x = plan.call(add, 2, 3)
registry.add(x, TestStore())
uberjob.render(uberjob.run(plan, registry=registry, dry_run=True),
format="svg")
uberjob.render(uberjob.run(plan, registry=registry, output=x,
dry_run=True),
format="svg")
def test_stack_frame_registry_source(self):
plan = uberjob.Plan()
registry = uberjob.Registry()
stack_frame = get_stack_frame(1)
x = registry.source(plan, TestStore())
with self.assert_call_exception(
expected_stack_frame=copy_with_line_offset(stack_frame, 1)):
uberjob.run(plan, output=x)
with self.assert_call_exception(
expected_stack_frame=copy_with_line_offset(stack_frame, 1)):
uberjob.run(plan, registry=registry, output=x)
def test_pruning(self):
p = uberjob.Plan()
r = uberjob.Registry()
x = p.call(operator.floordiv, 1, 0)
y = p.call(operator.floordiv, 14, 2)
self.assertEqual(uberjob.run(p, output=y), 7)
self.assertEqual(uberjob.run(p, output=y, registry=r), 7)
with self.assert_call_exception(ZeroDivisionError):
uberjob.run(p, output=x)
with self.assert_call_exception(ZeroDivisionError):
uberjob.run(p, output=x, registry=r)
def test_registry_type_validation(self):
plan = uberjob.Plan()
registry = uberjob.Registry()
a = plan.call(operator.add, 1, 2)
with self.assertRaises(TypeError):
registry.source(1, TestStore())
with self.assertRaises(TypeError):
registry.source(plan, 1)
with self.assertRaises(TypeError):
registry.add(1, TestStore())
with self.assertRaises(TypeError):
registry.add(a, 1)
def test_fresh_time_basic(self):
p = uberjob.Plan()
r = uberjob.Registry()
x = p.call(operator.add, 2, 3)
r.add(x, TestStore())
uberjob.run(p, registry=r)
self.assertEqual(r[x].write_count, 1)
uberjob.run(p, registry=r, fresh_time=r[x].modified_time)
self.assertEqual(r[x].write_count, 1)
uberjob.run(p,
registry=r,
fresh_time=r[x].modified_time + dt.timedelta(seconds=1))
self.assertEqual(r[x].write_count, 2)
def test_dry_run(self):
plan = uberjob.Plan()
registry = uberjob.Registry()
x = plan.lit(1)
self.assertEqual(
uberjob.run(plan, dry_run=True)[0].graph.number_of_nodes(), 0)
store = TestStore()
registry.add(x, store)
physical_plan, _ = uberjob.run(plan, registry=registry, dry_run=True)
self.assertEqual(plan.graph.number_of_nodes(), 1)
self.assertGreater(physical_plan.graph.number_of_nodes(), 1)
self.assertEqual(store.write_count, 0)
uberjob.run(physical_plan, output=list(physical_plan.graph.nodes()))
self.assertEqual(store.write_count, 1)
def test_registry_copy(self):
plan = uberjob.Plan()
registry = uberjob.Registry()
x = plan.call(operator.add, 2, 2)
y = plan.call(operator.add, 2, 2)
registry.add(x, TestStore())
registry_copy = registry.copy()
registry_copy.add(y, TestStore())
registry_copy.mapping[x].is_source = True
self.assertFalse(registry.mapping[x].is_source)
self.assertNotIn(y, registry)
self.assertIn(x, registry_copy)
self.assertEqual(registry[x], registry_copy[x])
def test_stale_source_successors_run_after_stale_source_predecessors(self):
p = uberjob.Plan()
r = uberjob.Registry()
t0 = dt.datetime.utcnow()
t1 = t0 + dt.timedelta(seconds=1)
x = r.source(p, TestStore(1, modified_time=t1))
s = TestStore(2, modified_time=t0)
y = p.call(s.write, x)
z = r.source(p, s)
p.add_dependency(y, z)
w = p.call(s.read)
p.add_dependency(z, w)
self.assertEqual(uberjob.run(p, registry=r, output=w), 1)
def test_dependent_source(self):
p = uberjob.Plan()
r = uberjob.Registry()
x = r.source(p, TestStore(0))
s = TestStore()
y = p.call(s.write, x)
z = r.source(p, s)
p.add_dependency(y, z)
self.assertEqual(uberjob.run(p, registry=r, output=z), 0)
self.assertEqual(s.read_count, 1)
self.assertEqual(s.write_count, 1)
self.assertEqual(uberjob.run(p, registry=r, output=z), 0)
self.assertEqual(s.read_count, 2)
self.assertEqual(s.write_count, 1)
def test_registry_simple(self):
p = uberjob.Plan()
r = uberjob.Registry()
x = p.call(operator.add, 2, 3)
s = TestStore()
r.add(x, s)
self.assertIn(x, r)
self.assertSetEqual({x}, set(r))
self.assertSetEqual({x}, set(r.keys()))
self.assertSetEqual({s}, set(r.values()))
self.assertSetEqual({(x, s)}, set(r.items()))
uberjob.run(p, registry=r)
self.assertEqual(r[x].read_count, 0)
self.assertEqual(r[x].write_count, 1)
uberjob.run(p, registry=r)
self.assertEqual(r[x].read_count, 0)
self.assertEqual(r[x].write_count, 1)
uberjob.run(p, registry=r, stale_check_max_workers=1)
def test_call_with_side_effects_but_no_args_or_return(self):
p = uberjob.Plan()
r = uberjob.Registry()
xs = TestStore(0)
x = r.source(p, xs)
ys = TestStore()
y = p.call(lambda: ys.write(xs.read()))
z = r.source(p, ys)
p.add_dependency(x, y)
p.add_dependency(y, z)
for i in range(2):
with self.subTest(i=i):
uberjob.run(p, registry=r)
self.assertEqual(xs.read_count, 1)
self.assertEqual(ys.read_count, 0)
self.assertEqual(ys.write_count, 1)
self.assertEqual(ys.value, 0)
def test_registry_complex(self):
for x_stale, y_stale, z_stale, z_output in itertools.product(
[False, True], repeat=4):
with self.subTest(x_stale=x_stale,
y_stale=y_stale,
z_stale=z_stale,
z_output=z_output):
p = uberjob.Plan()
r = uberjob.Registry()
x = p.call(operator.add, 2, 3)
y = p.call(operator.add, 4, 5)
z = p.call(operator.add, x, y)
r.add(
x,
TestStore(
Missing if x_stale else uberjob.run(p, output=x)))
r.add(
y,
TestStore(
Missing if y_stale else uberjob.run(p, output=y)))
r.add(
z,
TestStore(
Missing if z_stale else uberjob.run(p, output=z)))
self.assertEqual(
uberjob.run(p, output=z if z_output else None, registry=r),
14 if z_output else None,
)
self.assertEqual(r[x].read_count,
int(x_stale or y_stale or z_stale))
self.assertEqual(r[x].write_count, int(x_stale))
self.assertEqual(r[y].read_count,
int(x_stale or y_stale or z_stale))
self.assertEqual(r[y].write_count, int(y_stale))
self.assertEqual(r[z].read_count, int(z_output))
self.assertEqual(r[z].write_count,
int(x_stale or y_stale or z_stale))
def test_render_registry():
plan = uberjob.Plan()
registry = uberjob.Registry()
x = plan.call(add, 2, 3)
registry.add(x, TestStore())
uberjob.render(plan, registry=registry, format="svg")
def test_failed_to_read_from_empty_store_1(self):
p = uberjob.Plan()
r = uberjob.Registry()
x = r.source(p, TestStore())
with self.assert_failed_to_read_from_empty_store():
uberjob.run(p, registry=r, output=x)
def test_failed_to_get_modified_time(self):
p = uberjob.Plan()
r = uberjob.Registry()
x = r.source(p, TestStore(can_get_modified_time=False))
with self.assert_call_exception():
uberjob.run(p, registry=r, output=x, dry_run=True)
def test_forgot_registry_with_source(self):
p = uberjob.Plan()
r = uberjob.Registry()
x = r.source(p, TestStore(5))
with self.assert_forgotten_registry():
uberjob.run(p, output=x)