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