Implement OAuth refresh token flow

strawberryfields/api/connection.py

@@ -16,7 +16,7 @@
 """
 import io
 from datetime import datetime
-from typing import List
+from typing import List, Dict
 
 import numpy as np
 import requests
@@ -96,7 +96,8 @@ def __init__(
         self._verbose = verbose
 
         self._base_url = "http{}://{}:{}".format("s" if self.use_ssl else "", self.host, self.port)
-        self._headers = {"Authorization": self.token, "Accept-Version": self.api_version}
+
+        self._headers = {"Accept-Version": self.api_version}
 
         self.log = create_logger(__name__)
 
@@ -156,7 +157,7 @@ def get_device_spec(self, target: str) -> DeviceSpec:
     def _get_device_dict(self, target: str) -> dict:
         """Returns the device specifications as a dictionary"""
         path = f"/devices/{target}/specifications"
-        response = requests.get(self._url(path), headers=self._headers)
+        response = self._request("GET", self._url(path), headers=self._headers)
 
         if response.status_code == 200:
             self.log.info("The device spec %s has been successfully retrieved.", target)
@@ -185,7 +186,9 @@ def create_job(self, target: str, program: Program, run_options: dict = None) ->
         circuit = bb.serialize()
 
         path = "/jobs"
-        response = requests.post(self._url(path), headers=self._headers, json={"circuit": circuit})
+        response = self._request(
+            "POST", self._url(path), headers=self._headers, json={"circuit": circuit}
+        )
         if response.status_code == 201:
             job_id = response.json()["id"]
             if self._verbose:
@@ -221,7 +224,7 @@ def get_job(self, job_id: str) -> Job:
             strawberryfields.api.Job: the job
         """
         path = "/jobs/{}".format(job_id)
-        response = requests.get(self._url(path), headers=self._headers)
+        response = self._request("GET", self._url(path), headers=self._headers)
         if response.status_code == 200:
             return Job(
                 id_=response.json()["id"],
@@ -254,8 +257,8 @@ def get_job_result(self, job_id: str) -> Result:
             strawberryfields.api.Result: the job result
         """
         path = "/jobs/{}/result".format(job_id)
-        response = requests.get(
-            self._url(path), headers={"Accept": "application/x-numpy", **self._headers}
+        response = self._request(
+            "GET", self._url(path), headers={"Accept": "application/x-numpy", **self._headers}
         )
         if response.status_code == 200:
             # Read the numpy binary data in the payload into memory
@@ -283,8 +286,11 @@ def cancel_job(self, job_id: str):
             job_id (str): the job ID
         """
         path = "/jobs/{}".format(job_id)
-        response = requests.patch(
-            self._url(path), headers=self._headers, json={"status": JobStatus.CANCELLED.value}
+        response = self._request(
+            "PATCH",
+            self._url(path),
+            headers=self._headers,
+            json={"status": JobStatus.CANCELLED.value},
         )
         if response.status_code == 204:
             if self._verbose:
@@ -301,12 +307,56 @@ def ping(self) -> bool:
             bool: ``True`` if the connection is successful, and ``False`` otherwise
         """
         path = "/healthz"
-        response = requests.get(self._url(path), headers=self._headers)
+        response = self._request("GET", self._url(path), headers=self._headers)
         return response.status_code == 200
 
     def _url(self, path: str) -> str:
         return self._base_url + path
 
+    def _refresh_access_token(self):
+        """Use the offline token to request a new access token."""
+        self._headers.pop("Authorization", None)
+        path = "/auth/realms/platform/protocol/openid-connect/token"
+        headers = {**self._headers}
+        response = requests.post(
+            self._url(path),
+            headers=headers,
+            data={
+                "grant_type": "refresh_token",
+                "refresh_token": self._token,
+                "client_id": "public",
+            },
+        )
+        if response.status_code == 200:
+            access_token = response.json().get("access_token")
+            self._headers["Authorization"] = f"Bearer {access_token}"
+        else:
+            raise RequestFailedError(
+                "Authorization failed for request, please check your token provided."
+            )
+
+    def _request(self, method: str, path: str, headers: Dict = None, **kwargs):
+        """Wrap all API requests with an authentication token refresh if a 401 status
+        is received from the initial request.
+
+        Args:
+            method (str): the HTTP request method to use
+            path (str): path of the endpoint to use
+            headers (dict): dictionary containing the headers of the request
+
+        Returns:
+            requests.Response: the response received for the sent request
+        """
+        headers = headers or {}
+        request_headers = {**headers, **self._headers}
+        response = requests.request(method, path, headers=request_headers, **kwargs)
+        if response.status_code == 401:
+            # Refresh the access_token and retry the request
+            self._refresh_access_token()
+            request_headers = {**headers, **self._headers}
+            response = requests.request(method, path, headers=request_headers, **kwargs)
+        return response
+
     @staticmethod
     def _format_error_message(response: requests.Response) -> str:
         body = response.json()

tests/api/test_connection.py

@@ -50,6 +50,8 @@
 port = 443
 """
 
+test_host = "SomeHost"
+test_token = "SomeToken"
 
 class MockResponse:
     """A mock response with a JSON or binary body."""
@@ -72,9 +74,11 @@ def content(self):
 class TestConnection:
     """Tests for the ``Connection`` class."""
 
-    def test_init(self):
+    def test_init(self, monkeypatch):
         """Tests that a ``Connection`` is initialized correctly."""
         token, host, port, use_ssl = "token", "host", 123, True
+
+        monkeypatch.setattr(requests, "request", mock_return(MockResponse(200, {})))
         connection = Connection(token, host, port, use_ssl)
 
         assert connection.token == token
@@ -95,12 +99,13 @@ def test_get_device_spec(self, prog, connection, monkeypatch):
 
         monkeypatch.setattr(
             requests,
-            "get",
+            "request",
             mock_return(MockResponse(
                 200,
                 {"layout": "", "modes": 42, "compiler": [], "gate_parameters": {"param": [[0, 1]]}}
             )),
         )
+        monkeypatch.setattr(connection, "_refresh_access_token", lambda: None)
 
         device_spec = connection.get_device_spec(target)
 
@@ -114,7 +119,7 @@ def test_get_device_spec(self, prog, connection, monkeypatch):
 
     def test_get_device_spec_error(self, connection, monkeypatch):
         """Tests a failed device spec request."""
-        monkeypatch.setattr(requests, "get", mock_return(MockResponse(404, {})))
+        monkeypatch.setattr(requests, "request", mock_return(MockResponse(404, {})))
 
         with pytest.raises(RequestFailedError, match="Failed to get device specifications"):
             connection.get_device_spec("123")
@@ -124,7 +129,7 @@ def test_create_job(self, prog, connection, monkeypatch):
         id_, status = "123", JobStatus.QUEUED
 
         monkeypatch.setattr(
-            requests, "post", mock_return(MockResponse(201, {"id": id_, "status": status})),
+            requests, "request", mock_return(MockResponse(201, {"id": id_, "status": status})),
         )
 
         job = connection.create_job("X8_01", prog, {"shots": 1})
@@ -134,7 +139,7 @@ def test_create_job(self, prog, connection, monkeypatch):
 
     def test_create_job_error(self, prog, connection, monkeypatch):
         """Tests a failed job creation flow."""
-        monkeypatch.setattr(requests, "post", mock_return(MockResponse(400, {})))
+        monkeypatch.setattr(requests, "request", mock_return(MockResponse(400, {})))
 
         with pytest.raises(RequestFailedError, match="Failed to create job"):
             connection.create_job("X8_01", prog, {"shots": 1})
@@ -151,7 +156,7 @@ def test_get_all_jobs(self, connection, monkeypatch):
             for i in range(1, 10)
         ]
         monkeypatch.setattr(
-            requests, "get", mock_return(MockResponse(200, {"data": jobs})),
+            requests, "request", mock_return(MockResponse(200, {"data": jobs})),
         )
 
         jobs = connection.get_all_jobs(after=datetime(2020, 1, 5))
@@ -161,7 +166,7 @@ def test_get_all_jobs(self, connection, monkeypatch):
     @pytest.mark.xfail(reason="method not yet implemented")
     def test_get_all_jobs_error(self, connection, monkeypatch):
         """Tests a failed job list request."""
-        monkeypatch.setattr(requests, "get", mock_return(MockResponse(404, {})))
+        monkeypatch.setattr(requests, "request", mock_return(MockResponse(404, {})))
 
         with pytest.raises(RequestFailedError, match="Failed to get all jobs"):
             connection.get_all_jobs()
@@ -172,7 +177,7 @@ def test_get_job(self, connection, monkeypatch):
 
         monkeypatch.setattr(
             requests,
-            "get",
+            "request",
             mock_return(MockResponse(200, {"id": id_, "status": status.value, "meta": meta})),
         )
 
@@ -184,7 +189,7 @@ def test_get_job(self, connection, monkeypatch):
 
     def test_get_job_error(self, connection, monkeypatch):
         """Tests a failed job request."""
-        monkeypatch.setattr(requests, "get", mock_return(MockResponse(404, {})))
+        monkeypatch.setattr(requests, "request", mock_return(MockResponse(404, {})))
 
         with pytest.raises(RequestFailedError, match="Failed to get job"):
             connection.get_job("123")
@@ -195,15 +200,15 @@ def test_get_job_status(self, connection, monkeypatch):
 
         monkeypatch.setattr(
             requests,
-            "get",
+            "request",
             mock_return(MockResponse(200, {"id": id_, "status": status.value, "meta": {}})),
         )
 
         assert connection.get_job_status(id_) == status.value
 
     def test_get_job_status_error(self, connection, monkeypatch):
         """Tests a failed job status request."""
-        monkeypatch.setattr(requests, "get", mock_return(MockResponse(404, {})))
+        monkeypatch.setattr(requests, "request", mock_return(MockResponse(404, {})))
 
         with pytest.raises(RequestFailedError, match="Failed to get job"):
             connection.get_job_status("123")
@@ -231,7 +236,7 @@ def test_get_job_result(self, connection, result_dtype, monkeypatch):
             np.save(buf, result_samples)
             buf.seek(0)
             monkeypatch.setattr(
-                requests, "get", mock_return(MockResponse(200, binary_body=buf.getvalue())),
+                requests, "request", mock_return(MockResponse(200, binary_body=buf.getvalue())),
             )
 
         result = connection.get_job_result("123")
@@ -240,7 +245,7 @@ def test_get_job_result(self, connection, result_dtype, monkeypatch):
 
     def test_get_job_result_error(self, connection, monkeypatch):
         """Tests a failed job result request."""
-        monkeypatch.setattr(requests, "get", mock_return(MockResponse(404, {})))
+        monkeypatch.setattr(requests, "request", mock_return(MockResponse(404, {})))
 
         with pytest.raises(RequestFailedError, match="Failed to get job result"):
             connection.get_job_result("123")
@@ -255,30 +260,72 @@ def function(*args, **kwargs):
 
             return function
 
-        monkeypatch.setattr(requests, "patch", _mock_return(MockResponse(204, {})))
+        monkeypatch.setattr(requests, "request", _mock_return(MockResponse(204, {})))
 
         # A successful cancellation does not raise an exception
         connection.cancel_job("123")
 
     def test_cancel_job_error(self, connection, monkeypatch):
         """Tests a failed job cancellation request."""
-        monkeypatch.setattr(requests, "patch", mock_return(MockResponse(404, {})))
+        monkeypatch.setattr(requests, "request", mock_return(MockResponse(404, {})))
 
         with pytest.raises(RequestFailedError, match="Failed to cancel job"):
             connection.cancel_job("123")
 
     def test_ping_success(self, connection, monkeypatch):
         """Tests a successful ping to the remote host."""
-        monkeypatch.setattr(requests, "get", mock_return(MockResponse(200, {})))
+        monkeypatch.setattr(requests, "request", mock_return(MockResponse(200, {})))
 
         assert connection.ping()
 
     def test_ping_failure(self, connection, monkeypatch):
         """Tests a failed ping to the remote host."""
-        monkeypatch.setattr(requests, "get", mock_return(MockResponse(500, {})))
+        monkeypatch.setattr(requests, "request", mock_return(MockResponse(500, {})))
 
         assert not connection.ping()
 
+    def test_refresh_access_token(self, mocker, monkeypatch):
+        """Test that the access token is created by passing the expected headers."""
+        path = "/auth/realms/platform/protocol/openid-connect/token"
+
+        data={
+            "grant_type": "refresh_token",
+            "refresh_token": test_token,
+            "client_id": "public",
+        }
+
+        monkeypatch.setattr(requests, "post", mock_return(MockResponse(200, {})))
+        spy = mocker.spy(requests, "post")
+
+        conn = Connection(token=test_token, host=test_host)
+        conn._refresh_access_token()
+        expected_headers = {'Accept-Version': conn.api_version}
+        expected_url = f"https://{test_host}:443{path}"
+        spy.assert_called_once_with(expected_url, headers=expected_headers, data=data)
+
+    def test_refresh_access_token_raises(self, monkeypatch):
+        """Test that an error is raised when the access token could not be
+        generated while creating the Connection object."""
+        monkeypatch.setattr(requests, "post", mock_return(MockResponse(500, {})))
+        conn = Connection(token=test_token, host=test_host)
+        with pytest.raises(RequestFailedError, match="Authorization failed for request"):
+            conn._refresh_access_token()
+
+    def test_wrapped_request_refreshes(self, mocker, monkeypatch):
+        """Test that the _request method refreshes the access token when
+        getting a 401 response."""
+        # Mock post function used while refreshing
+        monkeypatch.setattr(requests, "post", mock_return(MockResponse(200, {})))
+
+        # Mock request function used for general requests
+        monkeypatch.setattr(requests, "request", mock_return(MockResponse(401, {})))
+
+        conn = Connection(token=test_token, host=test_host)
+
+        spy = mocker.spy(conn, "_refresh_access_token")
+        conn._request("SomeRequestMethod", "SomePath")
+        spy.assert_called_once_with()
+
 
 class TestConnectionIntegration:
     """Integration tests for using instances of the Connection."""