def test_secret_handler_can_use_any_secret_type(self):
class MySecret(prefect.tasks.secrets.Secret):
def run(self):
return "boo"
handler = SecretResultHandler(MySecret(name="foo"))
assert handler.write(123089123) == "foo"
assert handler.read(lambda x: None) == "boo"
def __init__(self, name, **kwargs):
kwargs["name"] = name
kwargs.setdefault("max_retries", 2)
kwargs.setdefault("retry_delay", datetime.timedelta(seconds=1))
if kwargs.get("result_handler"):
raise ValueError("Result Handlers for Secrets are not configurable.")
kwargs["result_handler"] = SecretResultHandler(secret_task=self)
super().__init__(**kwargs)
def test_task_failure_with_upstream_secrets_doesnt_store_secret_value_and_recompute_if_necessary(
client, ):
@prefect.task(max_retries=2, retry_delay=timedelta(seconds=100))
def is_p_three(p):
if p == 3:
raise ValueError("No thank you.")
return p
with prefect.Flow("test", result_handler=JSONResultHandler()) as f:
p = prefect.tasks.secrets.Secret("p")
res = is_p_three(p)
with prefect.context(secrets=dict(p=3)):
state = CloudFlowRunner(flow=f).run(return_tasks=[res])
assert state.is_running()
assert isinstance(state.result[res], Retrying)
exp_res = Result(3, result_handler=SecretResultHandler(p))
assert not state.result[res].cached_inputs["p"] == exp_res
exp_res.store_safe_value()
assert state.result[res].cached_inputs["p"] == exp_res
## here we set the result of the secret to a saferesult, ensuring
## it will get converted to a "true" result;
## we expect that the upstream value will actually get recomputed from context
## through the SecretResultHandler
safe = SafeResult("p", result_handler=SecretResultHandler(p))
state.result[p] = Success(result=safe)
state.result[res].start_time = pendulum.now("utc")
state.result[res].cached_inputs = dict(p=safe)
with prefect.context(secrets=dict(p=4)):
new_state = CloudFlowRunner(flow=f).run(return_tasks=[res],
task_states=state.result)
assert new_state.is_successful()
assert new_state.result[res].result == 4
def test_secret_handler_is_pickleable(self, secret_task):
handler = SecretResultHandler(secret_task)
new = cloudpickle.loads(cloudpickle.dumps(handler))
assert isinstance(new, SecretResultHandler)
def test_secret_handler_writes_and_reads(self, res, secret_task):
handler = SecretResultHandler(secret_task)
with set_temporary_config({"cloud.use_local_secrets": True}):
with prefect.context(secrets=dict(test=res)):
final = handler.read(handler.write(res))
assert final == res
def test_secret_handler_writes_by_only_returning_name(
self, res, secret_task):
handler = SecretResultHandler(secret_task)
out = handler.write(res)
assert out == "test"
def test_secret_handler_initializes_with_secret_task(self, secret_task):
handler = SecretResultHandler(secret_task=secret_task)
assert isinstance(handler.secret_task, prefect.tasks.secrets.Secret)
assert handler.secret_task.name == "test"
def test_secret_handler_requires_secret_task_at_init(self):
with pytest.raises(TypeError, match="missing 1 required position"):
handler = SecretResultHandler()
def test_serialize(self, secret_task):
serialized = ResultHandlerSchema().dump(SecretResultHandler(secret_task))
assert isinstance(serialized, dict)
assert serialized["type"] == "SecretResultHandler"
assert serialized["name"] == "foo"
def __init__(self, **kwargs):
if kwargs.get("result_handler"):
raise ValueError(
"Result Handlers for Secrets are not configurable.")
kwargs["result_handler"] = SecretResultHandler(secret_task=self)
super().__init__(**kwargs)
def test_basic_conversion_secret_result():
task = PrefectSecret("foo")
result_handler = SecretResultHandler(task)
result = ResultHandlerResult.from_result_handler(result_handler)
assert isinstance(result, SecretResult)
assert result.secret_task is task
