def test_tokens_are_passed_to_post(self, monkeypatch): post = MagicMock() session = MagicMock() session.return_value.post = post monkeypatch.setattr("requests.Session", session) with set_temporary_config({"cloud.graphql": "http://my-cloud.foo"}): client = Client() client.post("/foo/bar", token="secret_token") assert post.called assert post.call_args[1]["headers"] == {"Authorization": "Bearer secret_token"}
def test_headers_are_passed_to_post(self, monkeypatch): post = MagicMock() 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() client.post("/foo/bar", headers={"x": "y", "Authorization": "z"}) assert post.called assert post.call_args[1]["headers"] == { "x": "y", "Authorization": "Bearer secret_token", "X-PREFECT-CORE-VERSION": str(prefect.__version__), }
def test_client_posts_raises_with_no_token(monkeypatch): post = MagicMock() monkeypatch.setattr("requests.post", post) with set_temporary_config( {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": None} ): client = Client() with pytest.raises(AuthorizationError) as exc: result = client.post("/foo/bar") assert "Client.login" in str(exc.value)
def test_client_posts_to_api_server(patch_post): post = patch_post(dict(success=True)) with set_temporary_config( {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"} ): client = Client() result = client.post("/foo/bar") assert result == {"success": True} assert post.called assert post.call_args[0][0] == "http://my-cloud.foo/foo/bar"
def test_client_posts_raises_with_no_token(monkeypatch): post = MagicMock() 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": None }): client = Client() with pytest.raises(AuthorizationError, match="Client.login"): result = client.post("/foo/bar")
def test_client_posts_to_graphql_server(monkeypatch): post = MagicMock(return_value=MagicMock(json=MagicMock(return_value=dict( success=True)))) monkeypatch.setattr("requests.post", post) with set_temporary_config({ "cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token" }): client = Client() result = client.post("/foo/bar") assert result == {"success": True} assert post.called assert post.call_args[0][0] == "http://my-cloud.foo/foo/bar"
def test_client_posts_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.post("/foo/bar") 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/foo/bar" assert client.token == "new-token"
class CloudResultHandler(ResultHandler): """ Hook for storing and retrieving task results from Prefect cloud storage. Args: - result_handler_service (str, optional): the location of the service which will further process and store the results; if not provided, will default to the value of `cloud.result_handler` in your config file """ def __init__(self, result_handler_service: str = None) -> None: self._client = None if result_handler_service is None: self.result_handler_service = config.cloud.result_handler else: self.result_handler_service = result_handler_service super().__init__() def _initialize_client(self) -> None: """ Helper method for ensuring that CloudHandlers which are initialized locally do not attempt to start a Client. This is important because CloudHandlers are currently attached to `Flow` objects which need to be serialized / deserialized independently of cloud settings. This will instantiate a Client upon the first call to (de)serialize. """ if self._client is None: self._client = Client() # type: ignore def read(self, uri: str) -> Any: """ Read a result from the given URI location. Args: - uri (str): the path to the location of a result Returns: - the deserialized result from the provided URI """ self._initialize_client() self.logger.debug("Starting to read result from {}...".format(uri)) res = self._client.get( # type: ignore "/", server=self.result_handler_service, **{"uri": uri}) try: return_val = cloudpickle.loads( base64.b64decode(res.get("result", ""))) except EOFError: return_val = None self.logger.debug("Finished reading result from {}...".format(uri)) return return_val def write(self, result: Any) -> str: """ Write the provided result to Prefect Cloud. Args: - result (Any): the result to store Returns: - str: the URI path to the result in Cloud storage """ self._initialize_client() binary_data = base64.b64encode(cloudpickle.dumps(result)).decode() self.logger.debug("Starting to upload result to {}...".format( self.result_handler_service)) res = self._client.post( # type: ignore "/", server=self.result_handler_service, **{"result": binary_data}) self.logger.debug("Finished uploading result to {}...".format( self.result_handler_service)) return res.get("uri", "")