def test_successful_idempotent_migration(self):
pipeline = Pipeline("demo")
pipeline.add_state(PipelineState("state1", {"attr1": True}), is_starting=True)
pipeline.add_state(PipelineState("state2", {"attr2": True}))
pipeline.add_state(PipelineState("state3", {"attr3": True}))
self.content_store.append_multiple(
[
# belongs to state1
{"attr1": True, "idem": 1},
{"attr1": True, "idem": 2},
{"attr1": True, "idem": 3},
# belongs to state2
{"attr2": True, "idem": 4},
# belongs to state3
{"attr3": True, "idem": 5},
],
"idem",
)
pipeline.perform_state_migrations(self.content_store)
# simulate a failed migration here
self.content_store.update_many(
filter_q={"attr1": True}, update_doc={"$unset": {"_state": ""}}
)
pipeline.perform_state_migrations(self.content_store)
assert self.actual_documents_without_id() == [
{"attr1": True, "idem": 1, "_state": "state1"},
{"attr1": True, "idem": 2, "_state": "state1"},
{"attr1": True, "idem": 3, "_state": "state1"},
{"attr2": True, "idem": 4, "_state": "state2"},
{"attr3": True, "idem": 5, "_state": "state3"},
]
def test_assert_missing_starting_state(self):
pipeline = Pipeline("not_dag")
pipeline.add_state(PipelineState("state1", {}))
pipeline.add_state(PipelineState("state2", {}))
pipeline.add_state(PipelineState("state3", {}))
pipeline.add_worker(self.create_demo_worker_instance("state1", {"state2"}))
pipeline.add_mod_view(
PipelineModView(
"2to3", "state2", [PipelineModViewToState("state3", "state3")], [], ""
)
)
self.assertRaisesRegex(
RuntimeError, "starting state", pipeline.assert_pipeline_is_dag
)
def test_successful_migration(self):
pipeline = Pipeline("demo")
pipeline.add_state(PipelineState("state1", {"attr1": True}), is_starting=True)
pipeline.add_state(PipelineState("state2", {"attr2": True}))
pipeline.add_state(PipelineState("state3", {"attr3": True}))
self.content_store.append_multiple(
[
# belongs to state1
{"attr1": True, "idem": 1},
{"attr1": True, "idem": 2},
{"attr1": True, "idem": 3},
# belongs to state2
{"attr2": True, "idem": 4},
# belongs to state3
{"attr3": True, "idem": 5},
],
"idem",
)
pipeline.perform_state_migrations(self.content_store)
assert self.actual_documents_without_id() == [
{"attr1": True, "idem": 1, "_state": "state1"},
{"attr1": True, "idem": 2, "_state": "state1"},
{"attr1": True, "idem": 3, "_state": "state1"},
{"attr2": True, "idem": 4, "_state": "state2"},
{"attr3": True, "idem": 5, "_state": "state3"},
]
def test_assert_tricky_missing_and_duplicate_doc_invalid(self):
pipeline = Pipeline("demo")
pipeline.add_state(PipelineState("state1", {"attr1": True}), is_starting=True)
pipeline.add_state(PipelineState("state2", {"attr2": True}))
pipeline.add_state(PipelineState("state3", {"attr3": True}))
self.content_store.append_multiple(
[
# belongs to state1
{"attr1": True, "idem": 1},
{"attr1": True, "idem": 2},
{"attr1": True, "idem": 3},
# belongs to no state
{"attr4": True, "idem": 4},
# belongs to state2 and state3
{"attr2": True, "attr3": True, "idem": 5},
],
"idem",
)
self.assertRaisesRegex(
RuntimeError,
"is going to have both",
pipeline.assert_valid_state_migrations,
self.content_store,
2,
)
def test_assert_missing_doc_invalid(self):
pipeline = Pipeline("demo")
pipeline.add_state(PipelineState("state1", {"attr1": True}), is_starting=True)
pipeline.add_state(PipelineState("state2", {"attr2": True}))
pipeline.add_state(PipelineState("state3", {"attr3": True}))
self.content_store.append_multiple(
[
# belongs to state1
{"attr1": True, "idem": 1},
{"attr1": True, "idem": 2},
{"attr1": True, "idem": 3},
# belongs to state2
{"attr2": True, "idem": 4},
# belongs no state
{"attr4": True, "idem": 5},
],
"idem",
)
self.assertRaisesRegex(
RuntimeError,
"Sum of all state docs does not match total count",
pipeline.assert_valid_state_migrations,
self.content_store,
2,
)
def test_assert_valid_migration(self):
pipeline = Pipeline("demo")
pipeline.add_state(PipelineState("state1", {"attr1": True}), is_starting=True)
pipeline.add_state(PipelineState("state2", {"attr2": True}))
pipeline.add_state(PipelineState("state3", {"attr3": True}))
self.content_store.append_multiple(
[
# belongs to state1
{"attr1": True, "idem": 1},
{"attr1": True, "idem": 2},
{"attr1": True, "idem": 3},
# belongs to state2
{"attr2": True, "idem": 4},
# belongs to state3
{"attr3": True, "idem": 5},
],
"idem",
)
pipeline.assert_valid_state_migrations(self.content_store, 2)
def test_assert_not_dag(self):
pipeline = Pipeline("not_dag")
pipeline.add_state(PipelineState("state1", {}), is_starting=True)
pipeline.add_state(PipelineState("state2", {}))
pipeline.add_state(PipelineState("state3", {}))
pipeline.add_state(PipelineState("state4", {}))
pipeline.add_state(PipelineState("state5", {}))
pipeline.add_worker(self.create_demo_worker_instance("state1", {"state2"}))
pipeline.add_worker(self.create_demo_worker_instance("state3", {"state2"}))
pipeline.add_worker(self.create_demo_worker_instance("state2", {"state4"}))
pipeline.add_worker(self.create_demo_worker_instance("state4", {"state3"}))
pipeline.add_worker(self.create_demo_worker_instance("state4", {"state5"}))
self.assertRaisesRegex(
RuntimeError, "has a cycle", pipeline.assert_pipeline_is_dag
)
def test_assert_dag_simple_pipeline(self):
pipeline = Pipeline("simple")
pipeline.add_state(PipelineState("state1", {}), is_starting=True)
pipeline.add_state(PipelineState("state2", {}))
pipeline.add_state(PipelineState("state3", {}))
pipeline.add_state(PipelineState("state4", {}))
pipeline.add_state(PipelineState("state5", {}))
pipeline.add_worker(self.create_demo_worker_instance("state1", {"state2"}))
pipeline.add_worker(self.create_demo_worker_instance("state2", {"state3"}))
pipeline.add_worker(self.create_demo_worker_instance("state2", {"state4"}))
pipeline.add_worker(self.create_demo_worker_instance("state3", {"state4"}))
pipeline.add_worker(self.create_demo_worker_instance("state4", {"state5"}))
pipeline.assert_pipeline_is_dag()
def test_assert_dag_straight_line_pipeline(self):
pipeline = Pipeline("straight_line")
pipeline.add_state(PipelineState("state1", {}), is_starting=True)
pipeline.add_state(PipelineState("state2", {}))
pipeline.add_state(PipelineState("state3", {}))
pipeline.add_worker(self.create_demo_worker_instance("state1", {"state2"}))
pipeline.add_mod_view(
PipelineModView(
"2to3", "state2", [PipelineModViewToState("state3", "state3")], [], ""
)
)
pipeline.assert_pipeline_is_dag()