def test_client_deploy_rejects_setting_active_schedules_for_flows_with_req_params(
active, monkeypatch):
post = MagicMock()
monkeypatch.setattr("requests.post", post)
with set_temporary_config({
"cloud.graphql": "http://my-cloud.foo",
"cloud.auth_token": "secret_token"
}):
client = Client()
flow = prefect.Flow(name="test", schedule=prefect.schedules.Schedule())
flow.add_task(prefect.Parameter("x", required=True))
with pytest.raises(ClientError) as exc:
result = client.deploy(flow,
project_name="my-default-project",
set_schedule_active=active)
assert (str(
exc.value
) == "Flows with required parameters can not be scheduled automatically.")
def test_login_to_client_sets_access_token(self, patch_post):
tenant_id = str(uuid.uuid4())
post = patch_post({
"data": {
"tenant": [{
"id": tenant_id
}],
"switchTenant": {
"accessToken": "ACCESS_TOKEN",
"expiresAt": "2100-01-01",
"refreshToken": "REFRESH_TOKEN",
},
}
})
client = Client()
assert client._access_token is None
assert client._refresh_token is None
client.login_to_tenant(tenant_id=tenant_id)
assert client._access_token == "ACCESS_TOKEN"
assert client._refresh_token == "REFRESH_TOKEN"
def test_login_uses_api_token(self, patch_post):
tenant_id = str(uuid.uuid4())
post = patch_post({
"data": {
"tenant": [{
"id": tenant_id
}],
"switchTenant": {
"accessToken": "ACCESS_TOKEN",
"expiresAt": "2100-01-01",
"refreshToken": "REFRESH_TOKEN",
},
}
})
client = Client(api_token="api")
client.login_to_tenant(tenant_id=tenant_id)
assert post.call_args[1]["headers"] == {
"Authorization": "Bearer api",
"X-PREFECT-CORE-VERSION": str(prefect.__version__),
}
def test_set_flow_run_state_with_error(patch_post):
response = {
"data": {
"set_flow_run_state": None
},
"errors": [{
"message": "something went wrong"
}],
}
post = patch_post(response)
with set_temporary_config({
"cloud.api": "http://my-cloud.foo",
"cloud.auth_token": "secret_token"
}):
client = Client()
with pytest.raises(ClientError, match="something went wrong"):
client.set_flow_run_state(flow_run_id="74-salt",
version=0,
state=Pending())
def test_client_delete_project_error(patch_post, monkeypatch):
patch_post(
{
"data": {
"project": {},
}
}
)
project_name = "my-default-project"
monkeypatch.setattr(
"prefect.client.Client.get_default_tenant_slug", MagicMock(return_value="tslug")
)
with set_temporary_config({"cloud.auth_token": "secret_token", "backend": "cloud"}):
client = Client()
with pytest.raises(ValueError, match="Project {} not found".format(project_name)):
client.delete_project(project_name=project_name)
def test_get_task_run_info_with_error(patch_post):
response = {
"data": {"get_or_create_task_run": None},
"errors": [{"message": "something went wrong"}],
}
patch_post(response)
with set_temporary_config(
{
"cloud.api": "http://my-cloud.foo",
"cloud.auth_token": "secret_token",
"backend": "cloud",
}
):
client = Client()
with pytest.raises(ClientError, match="something went wrong"):
client.get_task_run_info(
flow_run_id="74-salt", task_id="72-salt", map_index=None
)
def test_client_deploy_with_flow_that_cant_be_deserialized(patch_post):
patch_post({"data": {"project": [{"id": "proj-id"}]}})
with set_temporary_config({
"cloud.graphql": "http://my-cloud.foo",
"cloud.auth_token": "secret_token"
}):
client = Client()
task = prefect.Task()
# we add a max_retries value to the task without a corresponding retry_delay; this will fail at deserialization
task.max_retries = 3
flow = prefect.Flow(name="test", tasks=[task])
with pytest.raises(
ValueError,
match=("(Flow could not be deserialized).*"
"(`retry_delay` must be provided if max_retries > 0)"),
) as exc:
client.deploy(flow, project_name="my-default-project", build=False)
def test_client_graphql_retries_if_token_needs_refreshing(monkeypatch):
error = requests.HTTPError()
error.response = MagicMock(status_code=401) # unauthorized
post = MagicMock(
return_value=MagicMock(
raise_for_status=MagicMock(side_effect=error),
json=MagicMock(return_value=dict(token="new-token")),
)
)
monkeypatch.setattr("requests.post", post)
with set_temporary_config(
{"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
):
client = Client()
with pytest.raises(requests.HTTPError) as exc:
result = client.graphql("{}")
assert exc.value is error
assert post.call_count == 3 # first call -> refresh token -> last call
assert post.call_args[0][0] == "http://my-cloud.foo"
assert client.token == "new-token"
def test_get_task_run_info(patch_post):
response = {
"getOrCreateTaskRun": {
"task_run": {
"id": "772bd9ee-40d7-479c-9839-4ab3a793cabd",
"version": 0,
"serialized_state": {
"type": "Pending",
"_result": {
"type": "SafeResult",
"value": "42",
"result_handler": {
"type": "JSONResultHandler"
},
},
"message": None,
"__version__": "0.3.3+310.gd19b9b7.dirty",
"cached_inputs": None,
},
"task": {
"slug": "slug"
},
}
}
}
post = patch_post(dict(data=response))
with set_temporary_config({
"cloud.graphql": "http://my-cloud.foo",
"cloud.auth_token": "secret_token"
}):
client = Client()
result = client.get_task_run_info(flow_run_id="74-salt",
task_id="72-salt",
map_index=None)
assert isinstance(result, TaskRunInfoResult)
assert isinstance(result.state, Pending)
assert result.state.result == "42"
assert result.state.message is None
assert result.id == "772bd9ee-40d7-479c-9839-4ab3a793cabd"
assert result.version == 0
def test_client_register(patch_post, compressed):
if compressed:
response = {
"data": {
"project": [{
"id": "proj-id"
}],
"createFlowFromCompressedString": {
"id": "long-id"
},
}
}
else:
response = {
"data": {
"project": [{
"id": "proj-id"
}],
"createFlow": {
"id": "long-id"
}
}
}
patch_post(response)
with set_temporary_config({
"cloud.graphql": "http://my-cloud.foo",
"cloud.auth_token": "secret_token"
}):
client = Client()
flow = prefect.Flow(name="test",
storage=prefect.environments.storage.Memory())
flow.result_handler = flow.storage.result_handler
flow_id = client.register(
flow,
project_name="my-default-project",
compressed=compressed,
version_group_id=str(uuid.uuid4()),
)
assert flow_id == "long-id"
def test_client_register_builds_flow(patch_post, compressed, monkeypatch):
if compressed:
response = {
"data": {
"project": [{"id": "proj-id"}],
"createFlowFromCompressedString": {"id": "long-id"},
}
}
else:
response = {
"data": {"project": [{"id": "proj-id"}], "createFlow": {"id": "long-id"}}
}
post = patch_post(response)
monkeypatch.setattr(
"prefect.client.Client.get_default_tenant_slug", MagicMock(return_value="tslug")
)
with set_temporary_config(
{"cloud.api": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
):
client = Client()
flow = prefect.Flow(name="test", storage=prefect.environments.storage.Memory())
flow.result_handler = flow.storage.result_handler
flow_id = client.register(
flow, project_name="my-default-project", compressed=compressed
)
## extract POST info
if compressed:
serialized_flow = decompress(
json.loads(post.call_args[1]["json"]["variables"])["input"][
"serializedFlow"
]
)
else:
serialized_flow = json.loads(post.call_args[1]["json"]["variables"])["input"][
"serializedFlow"
]
assert serialized_flow["storage"] is not None
def test_client_register_flow_id_output(
patch_post, compressed, monkeypatch, capsys, cloud_api, tmpdir
):
if compressed:
response = {
"data": {
"project": [{"id": "proj-id"}],
"create_flow_from_compressed_string": {"id": "long-id"},
}
}
else:
response = {
"data": {"project": [{"id": "proj-id"}], "create_flow": {"id": "long-id"}}
}
patch_post(response)
monkeypatch.setattr(
"prefect.client.Client.get_default_tenant_slug", MagicMock(return_value="tslug")
)
with set_temporary_config(
{
"cloud.api": "http://my-cloud.foo",
"cloud.auth_token": "secret_token",
"backend": "cloud",
}
):
client = Client()
flow = prefect.Flow(name="test", storage=prefect.environments.storage.Local(tmpdir))
flow.result = flow.storage.result
flow_id = client.register(
flow,
project_name="my-default-project",
compressed=compressed,
version_group_id=str(uuid.uuid4()),
)
assert flow_id == "long-id"
captured = capsys.readouterr()
assert "Flow: https://cloud.prefect.io/tslug/flow/long-id\n" in captured.out
def test_client_register_with_bad_proj_name(patch_post, monkeypatch,
cloud_api):
patch_post({"data": {"project": []}})
monkeypatch.setattr("prefect.client.Client.get_default_tenant_slug",
MagicMock(return_value="tslug"))
with set_temporary_config({
"cloud.auth_token": "secret_token",
"backend": "cloud"
}):
client = Client()
flow = prefect.Flow(name="test")
flow.result = prefect.engine.result.Result()
with pytest.raises(ValueError) as exc:
flow_id = client.register(flow,
project_name="my-default-project",
no_url=True)
assert "not found" in str(exc.value)
assert "prefect create project 'my-default-project'" in str(exc.value)
def test_refresh_token_sets_attributes(self, patch_post):
patch_post({
"data": {
"refreshToken": {
"accessToken": "ACCESS_TOKEN",
"expiresAt": "2100-01-01",
"refreshToken": "REFRESH_TOKEN",
}
}
})
client = Client()
assert client._access_token is None
assert client._refresh_token is None
# add buffer because Windows doesn't compare milliseconds
assert client._access_token_expires_at < pendulum.now().add(seconds=1)
client._refresh_access_token()
assert client._access_token is "ACCESS_TOKEN"
assert client._refresh_token is "REFRESH_TOKEN"
assert client._access_token_expires_at > pendulum.now().add(
seconds=599)
def test_client_create_project_that_already_exists(patch_posts, monkeypatch):
patch_posts(
[
{
"errors": [
{"message": "Uniqueness violation.", "path": ["create_project"]}
],
"data": {"create_project": None},
},
{"data": {"project": [{"id": "proj-id"}]}},
]
)
monkeypatch.setattr(
"prefect.client.Client.get_default_tenant_slug", MagicMock(return_value="tslug")
)
with set_temporary_config({"cloud.auth_token": "secret_token", "backend": "cloud"}):
client = Client()
project_id = client.create_project(project_name="my-default-project")
assert project_id == "proj-id"
def test_get_cloud_url_different_regex(patch_post, cloud_api):
response = {
"data": {"user": [{"default_membership": {"tenant": {"slug": "tslug"}}}]}
}
patch_post(response)
with set_temporary_config(
{
"cloud.api": "http://api-hello.prefect.io",
"cloud.auth_token": "secret_token",
"backend": "cloud",
}
):
client = Client()
url = client.get_cloud_url(subdirectory="flow", id="id")
assert url == "http://hello.prefect.io/tslug/flow/id"
url = client.get_cloud_url(subdirectory="flow-run", id="id2")
assert url == "http://hello.prefect.io/tslug/flow-run/id2"
def test_get_task_run_info(monkeypatch):
response = """
{
"getOrCreateTaskRun": {
"task_run": {
"id": "772bd9ee-40d7-479c-9839-4ab3a793cabd",
"version": 0,
"serialized_state": {
"type": "Pending",
"_result": {"type": "SafeResult", "value": "42", "result_handler": {"type": "JSONResultHandler"}},
"message": null,
"__version__": "0.3.3+310.gd19b9b7.dirty",
"cached_inputs": null
},
"task": {
"slug": "slug"
}
}
}
}
"""
post = MagicMock(return_value=MagicMock(json=MagicMock(return_value=dict(
data=json.loads(response)))))
session = MagicMock()
session.return_value.post = post
monkeypatch.setattr("requests.Session", session)
with set_temporary_config({
"cloud.graphql": "http://my-cloud.foo",
"cloud.auth_token": "secret_token"
}):
client = Client()
result = client.get_task_run_info(flow_run_id="74-salt",
task_id="72-salt",
map_index=None)
assert isinstance(result, TaskRunInfoResult)
assert isinstance(result.state, Pending)
assert result.state.result == "42"
assert result.state.message is None
assert result.id == "772bd9ee-40d7-479c-9839-4ab3a793cabd"
assert result.version == 0
def test_client_register_doesnt_raise_if_no_keyed_edges(
patch_post, compressed, monkeypatch, tmpdir
):
if compressed:
response = {
"data": {
"project": [{"id": "proj-id"}],
"create_flow_from_compressed_string": {"id": "long-id"},
}
}
else:
response = {
"data": {"project": [{"id": "proj-id"}], "create_flow": {"id": "long-id"}}
}
patch_post(response)
monkeypatch.setattr(
"prefect.client.Client.get_default_tenant_slug", MagicMock(return_value="tslug")
)
with set_temporary_config(
{
"cloud.api": "http://my-cloud.foo",
"cloud.auth_token": "secret_token",
"backend": "cloud",
}
):
client = Client()
flow = prefect.Flow(name="test", storage=prefect.environments.storage.Local(tmpdir))
flow.result = None
flow_id = client.register(
flow,
project_name="my-default-project",
compressed=compressed,
version_group_id=str(uuid.uuid4()),
no_url=True,
)
assert flow_id == "long-id"
def test_set_flow_run_state_with_error(monkeypatch):
response = {
"data": {
"setFlowRunState": None
},
"errors": [{
"message": "something went wrong"
}],
}
post = MagicMock(return_value=MagicMock(json=MagicMock(
return_value=response)))
monkeypatch.setattr("requests.post", post)
with set_temporary_config({
"cloud.graphql": "http://my-cloud.foo",
"cloud.auth_token": "secret_token"
}):
client = Client()
with pytest.raises(ClientError) as exc:
client.set_flow_run_state(flow_run_id="74-salt",
version=0,
state=Pending())
assert "something went wrong" in str(exc.value)
def test_client_deploy(monkeypatch):
response = {
"data": {
"project": [{
"id": "proj-id"
}],
"createFlow": {
"id": "long-id"
}
}
}
post = MagicMock(return_value=MagicMock(json=MagicMock(
return_value=response)))
monkeypatch.setattr("requests.post", post)
with set_temporary_config({
"cloud.graphql": "http://my-cloud.foo",
"cloud.auth_token": "secret_token"
}):
client = Client()
flow = prefect.Flow(name="test")
flow_id = client.deploy(flow, project_name="my-default-project")
assert flow_id == "long-id"
def test_artifacts_client_functions(patch_post, cloud_api):
response = {
"data": {
"create_task_run_artifact": {"id": "artifact_id"},
"update_task_run_artifact": {"success": True},
"delete_task_run_artifact": {"success": True},
}
}
patch_post(response)
client = Client()
artifact_id = client.create_task_run_artifact(
task_run_id="tr_id", kind="kind", data={"test": "data"}, tenant_id="t_id"
)
assert artifact_id == "artifact_id"
client.update_task_run_artifact(task_run_artifact_id="tra_id", data={"new": "data"})
client.delete_task_run_artifact(task_run_artifact_id="tra_id")
response = {
"data": {
"create_task_run_artifact": {"id": None},
}
}
patch_post(response)
with pytest.raises(ValueError):
client.create_task_run_artifact(
task_run_id="tr_id", kind="kind", data={"test": "data"}, tenant_id="t_id"
)
with pytest.raises(ValueError):
client.update_task_run_artifact(task_run_artifact_id=None, data={"new": "data"})
with pytest.raises(ValueError):
client.delete_task_run_artifact(task_run_artifact_id=None)
def test_get_available_tenants(self, patch_post):
tenants = [
{
"id": "a",
"name": "a-name",
"slug": "a-slug"
},
{
"id": "b",
"name": "b-name",
"slug": "b-slug"
},
{
"id": "c",
"name": "c-name",
"slug": "c-slug"
},
]
post = patch_post({"data": {"tenant": tenants}})
client = Client()
gql_tenants = client.get_available_tenants()
assert gql_tenants == tenants
def test_client_deploy(monkeypatch, compressed):
if compressed:
response = {
"data": {
"project": [{
"id": "proj-id"
}],
"createFlowFromCompressedString": {
"id": "long-id"
},
}
}
else:
response = {
"data": {
"project": [{
"id": "proj-id"
}],
"createFlow": {
"id": "long-id"
}
}
}
post = MagicMock(return_value=MagicMock(json=MagicMock(
return_value=response)))
session = MagicMock()
session.return_value.post = post
monkeypatch.setattr("requests.Session", session)
with set_temporary_config({
"cloud.graphql": "http://my-cloud.foo",
"cloud.auth_token": "secret_token"
}):
client = Client()
flow = prefect.Flow(name="test",
storage=prefect.environments.storage.Memory())
flow_id = client.deploy(flow,
project_name="my-default-project",
compressed=compressed)
assert flow_id == "long-id"
def get(self) -> Optional[Any]:
"""
Retrieve the secret value. If not found, returns `None`.
If using local secrets, `Secret.get()` will attempt to call `json.loads` on the
value pulled from context. For this reason it is recommended to store local secrets as
JSON documents to avoid ambiguous behavior.
Returns:
- Any: the value of the secret; if not found, returns `None`
Raises:
- ValueError: if `.get()` is called within a Flow building context
- ValueError: if `use_local_secrets=False` and the Client fails to retrieve your secret
"""
if isinstance(prefect.context.get("flow"), prefect.core.flow.Flow):
raise ValueError(
"Secrets should only be retrieved during a Flow run, not while building a Flow."
)
if prefect.config.cloud.use_local_secrets is True:
secrets = prefect.context.get("secrets", {})
value = secrets.get(self.name)
try:
return json.loads(value)
except (json.JSONDecodeError, TypeError):
return value
else:
client = Client()
result = client.graphql(
"""
query($name: String!) {
secretValue(name: $name)
}
""",
name=self.name,
) # type: Any
return as_nested_dict(result.data.secretValue, dict)
def test_write_log_with_error(monkeypatch):
response = {
"data": {
"writeRunLog": None
},
"errors": [{
"message": "something went wrong"
}],
}
post = MagicMock(return_value=MagicMock(json=MagicMock(
return_value=response)))
session = MagicMock()
session.return_value.post = post
monkeypatch.setattr("requests.Session", session)
with set_temporary_config({
"cloud.graphql": "http://my-cloud.foo",
"cloud.auth_token": "secret_token"
}):
client = Client()
with pytest.raises(ClientError, match="something went wrong"):
client.write_run_log(flow_run_id="1")
def test_set_task_run_state_serializes(patch_post):
response = {
"data": {
"set_task_run_states": {
"states": [{
"status": "SUCCESS"
}]
}
}
}
post = patch_post(response)
with set_temporary_config({
"cloud.api": "http://my-cloud.foo",
"cloud.auth_token": "secret_token"
}):
client = Client()
res = SafeResult(lambda: None, result_handler=None)
with pytest.raises(marshmallow.exceptions.ValidationError):
client.set_task_run_state(task_run_id="76-salt",
version=0,
state=Pending(result=res))
def test_set_task_run_state(patch_post):
response = {
"data": {
"set_task_run_states": {
"states": [{
"status": "SUCCESS"
}]
}
}
}
post = patch_post(response)
state = Pending()
with set_temporary_config({
"cloud.api": "http://my-cloud.foo",
"cloud.auth_token": "secret_token"
}):
client = Client()
result = client.set_task_run_state(task_run_id="76-salt",
version=0,
state=state)
assert result is state
def test_get_default_tenant_slug_as_user(patch_post):
response = {
"data": {
"user": [{
"default_membership": {
"tenant": {
"slug": "tslug"
}
}
}]
}
}
patch_post(response)
with set_temporary_config({
"cloud.api": "http://my-cloud.foo",
"cloud.auth_token": "secret_token"
}):
client = Client()
slug = client.get_default_tenant_slug()
assert slug == "tslug"
def test_set_flow_run_state(patch_post):
response = {
"data": {
"set_flow_run_states": {
"states": [{"id": 1, "status": "SUCCESS", "message": None}]
}
}
}
patch_post(response)
with set_temporary_config(
{
"cloud.api": "http://my-cloud.foo",
"cloud.auth_token": "secret_token",
"backend": "cloud",
}
):
client = Client()
state = Pending()
result = client.set_flow_run_state(flow_run_id="74-salt", version=0, state=state)
assert isinstance(result, State)
assert isinstance(result, Pending)
def test_set_flow_run_state_gets_queued(patch_post):
response = {
"data": {
"set_flow_run_states": {
"states": [{"id": "74-salt", "status": "QUEUED", "message": None}]
}
}
}
post = patch_post(response)
with set_temporary_config(
{
"cloud.api": "http://my-cloud.foo",
"cloud.auth_token": "secret_token",
"backend": "cloud",
}
):
client = Client()
state = Running()
result = client.set_flow_run_state(flow_run_id="74-salt", version=0, state=state)
assert isinstance(result, State)
assert state != result
assert result.is_queued()
