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()