def client(monkeypatch):
cloud_client = MagicMock(
graphql=MagicMock(return_value=MagicMock(data=MagicMock(
flow=[MagicMock(
id="abc123"), MagicMock(id="def456")]))),
create_flow_run=MagicMock(return_value="xyz890"),
get_cloud_url=MagicMock(
return_value="https://api.prefect.io/flow/run/url"),
create_task_run_artifact=MagicMock(return_value="id"),
get_flow_run_info=MagicMock(return_value=FlowRunInfoResult(
id="my-flow-run-id",
name="test-run",
flow_id="xyz890",
version=1,
task_runs=[],
state=state.Success(),
scheduled_start_time=None,
project=ProjectInfo(id="my-project-id", name="Test Project"),
parameters={"test": "ing"},
context={},
)),
)
monkeypatch.setattr("prefect.tasks.prefect.flow_run.Client",
MagicMock(return_value=cloud_client))
monkeypatch.setattr("prefect.Client", MagicMock(return_value=cloud_client))
yield cloud_client
def test_new_result_with_no_location_serializes_correctly(self):
s = state.Success(message="test",
result=results.S3Result(bucket="foo"))
serialized = StateSchema().dump(s)
assert serialized["message"] == "test"
assert serialized["_result"]["type"] == "S3Result"
assert serialized["_result"]["location"] is None
def test_new_result_with_location_serializes_correctly(self):
s = state.Success(
message="test",
result=results.S3Result(bucket="foo", location="dir/place.txt"),
)
serialized = StateSchema().dump(s)
assert serialized["message"] == "test"
assert serialized["_result"]["type"] == "S3Result"
def test_deserialize_mapped():
s = state.Success(message="1", result=1)
f = state.Failed(message="2", result=2)
serialized = StateSchema().dump(state.Mapped(message="message", map_states=[s, f]))
deserialized = StateSchema().load(serialized)
assert isinstance(deserialized, state.Mapped)
assert len(deserialized.map_states) == 2
assert deserialized.map_states == [None, None]
assert deserialized._result == NoResult
def test_result_raises_error_on_dump_if_not_valid_json():
res = SafeResult({"x": {
"y": {
"z": lambda: 1
}
}},
result_handler=JSONResultHandler())
s = state.Success(result=res)
with pytest.raises(marshmallow.exceptions.ValidationError):
StateSchema().dump(s)
def test_result_must_be_valid_json():
res = SafeResult({"x": {
"y": {
"z": 1
}
}},
result_handler=JSONResultHandler())
s = state.Success(result=res)
serialized = StateSchema().dump(s)
assert serialized["_result"]["value"] == s.result
def test_state():
"""
This test ensures that deserialization works even when
another Marshmallow serializer called "StateSchema" has been
created.
https://github.com/PrefectHQ/prefect/pull/2738
"""
a = state.Submitted(state=state.Success())
prefect.serialization.state.StateSchema().load(a.serialize())
def test_serialize_mapped():
s = state.Success(message="1", result=1)
f = state.Failed(message="2", result=2)
serialized = StateSchema().dump(state.Mapped(message="message", map_states=[s, f]))
assert isinstance(serialized, dict)
assert serialized["type"] == "Mapped"
assert serialized["message"] == "message"
assert "_result" not in serialized
assert "map_states" not in serialized
assert serialized["n_map_states"] == 2
assert serialized["__version__"] == prefect.__version__
def test_new_result_with_location_deserializes_correctly(self):
s = state.Success(
message="test",
result=results.S3Result(bucket="foo", location="dir/place.txt"),
)
schema = StateSchema()
new_state = schema.load(schema.dump(s))
assert new_state.is_successful()
assert new_state.result is None
assert new_state._result.bucket == "foo"
assert isinstance(new_state._result, results.S3Result)
assert new_state._result.location == "dir/place.txt"
def test_raises_success_state(self, mock_watch_flow_run, MockFlowRunView):
MockFlowRunView.from_flow_run_id().get_latest().state = state.Success(
message="foo")
with prefect.Flow("test") as flow:
ref = wait_for_flow_run("flow-run-id", raise_final_state=True)
flow_state = flow.run()
task_state = flow_state.result[ref]
assert task_state.is_successful()
assert task_state.message == 'flow-run-id finished in state <Success: "foo">'
# The latest view is attached to the result
assert task_state.result == MockFlowRunView.from_flow_run_id(
).get_latest()
assert flow_state.is_successful()
class TestInitializeRun:
def test_initialize_run_returns_state_and_context(self):
state_, context = state.Pending(), {}
s, c = Runner().initialize_run(state=state_, context=context)
assert s is state_
assert c is context
@pytest.mark.parametrize(
"state",
[
state.Success(),
state.Failed(),
state.Pending(),
state.Scheduled(),
state.Skipped(),
state.Cached(),
state.Retrying(),
state.Running(),
],
)
def test_initialize_run_returns_state(self, state):
new_state, _ = Runner().initialize_run(state, context={})
assert new_state is state
@pytest.mark.parametrize(
"state",
[
state.Submitted(state=state.Pending()),
state.Submitted(state=state.Retrying()),
state.Submitted(state=state.Scheduled()),
state.Submitted(state=state.Resume()),
state.Queued(state=state.Pending()),
state.Queued(state=state.Retrying()),
state.Queued(state=state.Scheduled()),
state.Queued(state=state.Resume()),
],
)
def test_initialize_run_gets_wrapped_state_from_submitted_states(
self, state):
new_state, _ = Runner().initialize_run(state, context={})
assert new_state is state.state
def test_initialize_run_creates_pending_if_no_state_provided(self):
new_state, _ = Runner().initialize_run(state=None, context={})
assert isinstance(new_state, state.Pending)
def test_new_result_with_location_deserializes_correctly(self):
s = state.Success(
message="test",
result=results.S3Result(bucket="foo", location="dir/place.txt"),
)
def test_new_result_with_no_location_serializes_as_no_result(self):
s = state.Success(message="test", result=results.S3Result(bucket="foo"))
def test_deserialize_mapped():
s = state.Success(message="1", result=1)
with pytest.raises(marshmallow.exceptions.ValidationError):
StateSchema().load({"type": "FakeState"})
@pytest.mark.parametrize("state", complex_states())
def test_complex_state_attributes_are_handled(state):
serialized = StateSchema().dump(state)
deserialized = StateSchema().load(serialized)
assert state == deserialized
<<<<<<< HEAD
=======
def test_result_must_be_valid_json():
res = SafeResult({"x": {"y": {"z": 1}}}, result_handler=JSONResultHandler())
s = state.Success(result=res)
serialized = StateSchema().dump(s)
assert serialized["_result"]["value"] == s.result
def test_result_raises_error_on_dump_if_not_valid_json():
res = SafeResult({"x": {"y": {"z": lambda: 1}}}, result_handler=JSONResultHandler())
s = state.Success(result=res)
with pytest.raises(marshmallow.exceptions.ValidationError):
StateSchema().dump(s)
>>>>>>> prefect clone
def test_deserialize_json_with_context():
deserialized = StateSchema().load(
{"type": "Running", "context": {"boo": ["a", "b", "c"]}}
def run(self):
sleep(2)
return state.Success()