Repository: dwavesystems/dwave-cloud-client
Pull request #584 (2023-10-01 00:18:03)

Improve auth flow
In this PR we:
- make Ocean client id external configurable (via new config option `leap_client_id`)
- fix `LeapAuthFlow` requirements on endpoints available in config (they are now auto-resolved)
- fix caching of regions metadata to be persistent over program runs, also increase default max age to 7 days
- expose/proxy token handling methods on `AuthFlow`


Patch:
diff --git a/dwave/cloud/auth/config.py b/dwave/cloud/auth/config.py
index a8b0ea9..555e8a4 100644
--- a/dwave/cloud/auth/config.py
+++ b/dwave/cloud/auth/config.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 
 # XXX: update before releasing 0.11.0
-# TODO: make it externally configurable
 OCEAN_SDK_CLIENT_ID = '968678'
 
 OCEAN_SDK_SCOPES = (
diff --git a/dwave/cloud/auth/flows.py b/dwave/cloud/auth/flows.py
index c428367..90b6b7d 100644
--- a/dwave/cloud/auth/flows.py
+++ b/dwave/cloud/auth/flows.py
@@ -26,6 +26,7 @@ from authlib.common.security import generate_token
 from dwave.cloud.auth.config import OCEAN_SDK_CLIENT_ID, OCEAN_SDK_SCOPES
 from dwave.cloud.auth.server import SingleRequestAppServer, RequestCaptureApp
 from dwave.cloud.config.models import ClientConfig
+from dwave.cloud.regions import resolve_endpoints
 from dwave.cloud.utils import pretty_argvalues
 
 __all__ = ['AuthFlow', 'LeapAuthFlow']
@@ -75,7 +76,10 @@ class AuthFlow:
 
         self.session = OAuth2Session(
             client_id=client_id, scope=scopes, redirect_uri=redirect_uri,
-            code_challenge_method='S256')
+            code_challenge_method='S256',
+            # metadata set via kwargs
+            authorization_endpoint=authorization_endpoint,
+            token_endpoint=token_endpoint)
 
         if session_config is not None:
             self.update_session(session_config)
@@ -103,6 +107,14 @@ class AuthFlow:
     def redirect_uri(self, value):
         self.session.redirect_uri = value
 
+    @property
+    def token(self):
+        return self.session.token
+
+    @token.setter
+    def token(self, value):
+        self.session.token = value
+
     def get_authorization_url(self) -> str:
         self.state = generate_token(30)
         self.code_verifier = generate_token(48)
@@ -145,6 +157,12 @@ class AuthFlow:
         logger.debug(f"{type(self).__name__}.fetch_token() = {token!r}")
         return token
 
+    def refresh_token(self):
+        return self.session.refresh_token(url=self.token_endpoint)
+
+    def ensure_active_token(self):
+        return self.session.ensure_active_token(token=self.session.token)
+
 
 class LeapAuthFlow(AuthFlow):
     """:class:`.AuthFlow` specialized for Ocean and Leap.
@@ -203,6 +221,8 @@ class LeapAuthFlow(AuthFlow):
         logger.trace(f"{cls.__name__}.from_config_model("
                      f"config={config!r}, **kwargs={kwargs!r})")
 
+        config = resolve_endpoints(config, inplace=False)
+
         authorization_endpoint = cls._infer_auth_endpoint(config.leap_api_endpoint)
         token_endpoint = cls._infer_token_endpoint(config.leap_api_endpoint)
 
@@ -216,7 +236,7 @@ class LeapAuthFlow(AuthFlow):
             verify=not config.permissive_ssl)
 
         return cls(
-            client_id=kwargs.pop('client_id', OCEAN_SDK_CLIENT_ID),
+            client_id=kwargs.pop('client_id', config.leap_client_id or OCEAN_SDK_CLIENT_ID),
             scopes=kwargs.pop('scopes', OCEAN_SDK_SCOPES),
             redirect_uri=kwargs.pop('redirect_uri', cls._OOB_REDIRECT_URI),
             authorization_endpoint=authorization_endpoint,
diff --git a/dwave/cloud/client/base.py b/dwave/cloud/client/base.py
index 9ec7691..2fc0909 100644
--- a/dwave/cloud/client/base.py
+++ b/dwave/cloud/client/base.py
@@ -224,6 +224,10 @@ class Client(object):
         leap_api_endpoint (str, optional):
             Leap API endpoint.
 
+        leap_client_id (str, optional):
+            Leap OAuth 2.0 Ocean client id. Reserved for testing, otherwise
+            don't override.
+
         defaults (dict, optional):
             Defaults for the client instance that override the class
             :attr:`.Client.DEFAULTS`.
@@ -232,6 +236,9 @@ class Client(object):
         Added the ``leap_api_endpoint`` parameter and config option (also
         available via environment variable ``DWAVE_LEAP_API_ENDPOINT``).
 
+        Added the ``leap_client_id`` parameter and config option (reserved for
+        testing).
+
     Note:
         Default values of all constructor arguments listed above are kept in
         a class variable :attr:`.Client.DEFAULTS`.
@@ -285,6 +292,7 @@ class Client(object):
         'metadata_api_endpoint': api.constants.DEFAULT_METADATA_API_ENDPOINT,
         'region': DEFAULT_API_REGION,
         'leap_api_endpoint': None,
+        'leap_client_id': None,
         'endpoint': None,       # defined via region, resolved on client init
         'token': None,
         'solver': None,
@@ -333,7 +341,7 @@ class Client(object):
     _SOLVERS_CACHE_MAXAGE = 300     # 5 min
 
     # Downloaded region metadata cache maxage [sec]
-    _REGIONS_CACHE_MAXAGE = 86400   # 1 day
+    _REGIONS_CACHE_MAXAGE = 7 * 86400   # 7 days
 
     # Multipart upload parameters
     _UPLOAD_PART_SIZE_BYTES = 5 * 1024 * 1024
diff --git a/dwave/cloud/config/models.py b/dwave/cloud/config/models.py
index c346a7d..eed987f 100644
--- a/dwave/cloud/config/models.py
+++ b/dwave/cloud/config/models.py
@@ -117,6 +117,9 @@ class ClientConfig(BaseModel, GetterMixin):
     endpoint: Optional[str] = None
     token: Optional[str] = None
 
+    # oauth 2.0 support
+    leap_client_id: Optional[str] = None
+
     # [sapi client specific] feature-based solver selection query
     client: Optional[str] = None
     solver: Optional[Dict[str, Any]] = None
@@ -272,6 +275,7 @@ _V1_CONFIG_DEFAULTS = {
     'region': DEFAULT_REGION,
     'endpoint': None,
     'token': None,
+    'leap_client_id': None,
     'solver': None,
     'proxy': None,
     'permissive_ssl': False,
diff --git a/dwave/cloud/regions.py b/dwave/cloud/regions.py
index 6a1943c..1856712 100644
--- a/dwave/cloud/regions.py
+++ b/dwave/cloud/regions.py
@@ -22,7 +22,8 @@ from pydantic import TypeAdapter
 
 from dwave.cloud import api
 from dwave.cloud.api.constants import (
-    DEFAULT_REGION, DEFAULT_SOLVER_API_ENDPOINT, DEFAULT_LEAP_API_ENDPOINT)
+    DEFAULT_REGION, DEFAULT_SOLVER_API_ENDPOINT, DEFAULT_LEAP_API_ENDPOINT,
+    DEFAULT_METADATA_API_ENDPOINT)
 from dwave.cloud.config.models import ClientConfig, validate_config_v1
 from dwave.cloud.utils import cached
 
@@ -30,10 +31,10 @@ __all__ = ['get_regions']
 
 logger = logging.getLogger(__name__)
 
-_REGIONS_CACHE_MAXAGE = 86400   # 1 day
+_REGIONS_CACHE_MAXAGE = 7 * 86400   # 7 days
 
 
-@cached.ondisk(maxage=_REGIONS_CACHE_MAXAGE, key='cache_key')
+@cached.ondisk(maxage=_REGIONS_CACHE_MAXAGE, key='cache_key', bucket='regions')
 def _fetch_available_regions(cache_key: Any, config: ClientConfig) -> Dict[str, str]:
     logger.debug("Fetching available regions from the Metadata API at %r",
                  config.metadata_api_endpoint)
@@ -132,7 +133,17 @@ def _infer_leap_api_endpoint(solver_api_endpoint: str,
     return parts._replace(netloc=netloc, path=path).geturl()
 
 
-def resolve_endpoints(config: ClientConfig, *, inplace: bool = False) -> ClientConfig:
+def _infer_solver_api_endpoint(leap_api_endpoint: str) -> str:
+    # solver_api_endpoint from leap_api_endpoint
+
+    # sapi path
+    path = urlsplit(DEFAULT_SOLVER_API_ENDPOINT).path
+
+    return urlsplit(leap_api_endpoint)._replace(path=path).geturl()
+
+
+def resolve_endpoints(config: ClientConfig, *, inplace: bool = False,
+                      shortcircuit: bool = True) -> ClientConfig:
     """Use region and endpoint from configuration to resolve all endpoints.
 
     Explicit endpoint overrides the region (i.e. region extension is
@@ -142,9 +153,6 @@ def resolve_endpoints(config: ClientConfig, *, inplace: bool = False) -> ClientC
     available, default global endpoint is used.
     """
 
-    # Note: this function is currently indirectly tested with `Client`
-    # multi-region support tests.
-
     if not inplace:
         config = config.model_copy(deep=True)
 
@@ -153,6 +161,10 @@ def resolve_endpoints(config: ClientConfig, *, inplace: bool = False) -> ClientC
         config.endpoint = DEFAULT_SOLVER_API_ENDPOINT
         config.leap_api_endpoint = DEFAULT_LEAP_API_ENDPOINT
 
+    # we always need metadata endpoint
+    if config.metadata_api_endpoint is None:
+        config.metadata_api_endpoint = DEFAULT_METADATA_API_ENDPOINT
+
     # for backward-compat: endpoint overrides region
     if config.endpoint:
         if not config.leap_api_endpoint:
@@ -160,7 +172,19 @@ def resolve_endpoints(config: ClientConfig, *, inplace: bool = False) -> ClientC
                 solver_api_endpoint=config.endpoint, region_code=config.region)
         return config
 
+    # for consistency: any endpoint overrides region
+    if config.leap_api_endpoint:
+        if not config.endpoint:
+            config.endpoint = _infer_solver_api_endpoint(
+                leap_api_endpoint=config.leap_api_endpoint)
+        return config
+
     if not config.region:
+        config.region = DEFAULT_REGION
+
+    # short-circuit a Metadata API hit if we just need default values
+    if (shortcircuit and config.region == DEFAULT_REGION
+        and config.metadata_api_endpoint == DEFAULT_METADATA_API_ENDPOINT):
         set_defaults(config)
         return config
 
diff --git a/dwave/cloud/utils.py b/dwave/cloud/utils.py
index adda233..18d8f69 100644
--- a/dwave/cloud/utils.py
+++ b/dwave/cloud/utils.py
@@ -620,7 +620,7 @@ class cached:
             def f(x):
                 return x**2
 
-            @cached.disabled
+            @cached.disabled()
             def no_cache(x):
                 return f(x)
 
diff --git a/releasenotes/notes/add-leap-client-id-config-option-8c95422b741eda44.yaml b/releasenotes/notes/add-leap-client-id-config-option-8c95422b741eda44.yaml
new file mode 100644
index 0000000..a60fd3e
--- /dev/null
+++ b/releasenotes/notes/add-leap-client-id-config-option-8c95422b741eda44.yaml
@@ -0,0 +1,5 @@
+---
+features:
+  - |
+    Add a  configuration option for setting Leap OAuth 2.0 Ocean client id,
+    ``leap_client_id``. Primarily used for testing.


Test patch:
diff --git a/tests/auth/test_flows.py b/tests/auth/test_flows.py
index ebf3bd2..c33ced6 100644
--- a/tests/auth/test_flows.py
+++ b/tests/auth/test_flows.py
@@ -16,7 +16,7 @@ import threading
 import unittest
 from functools import partial
 from unittest import mock
-from urllib.parse import urlsplit, parse_qsl
+from urllib.parse import urlsplit, parse_qsl, urljoin
 
 import requests
 import requests_mock
@@ -24,6 +24,7 @@ import requests_mock
 from dwave.cloud.auth.flows import AuthFlow, LeapAuthFlow
 from dwave.cloud.auth.config import OCEAN_SDK_CLIENT_ID, OCEAN_SDK_SCOPES
 from dwave.cloud.config import ClientConfig
+from dwave.cloud.api.constants import DEFAULT_LEAP_API_ENDPOINT
 
 
 class TestAuthFlow(unittest.TestCase):
@@ -34,6 +35,7 @@ class TestAuthFlow(unittest.TestCase):
         self.redirect_uri_oob = 'oob'
         self.authorization_endpoint = 'https://example.com/authorize'
         self.token_endpoint = 'https://example.com/token'
+        self.token = dict(access_token='123', refresh_token='456', id_token='789')
 
         self.test_args = dict(
             client_id=self.client_id,
@@ -72,7 +74,6 @@ class TestAuthFlow(unittest.TestCase):
     def test_fetch_token(self, m):
         # mock the token_endpoint
         code = '123456'
-        token = dict(access_token='123', refresh_token='456', id_token='789')
         expected_params = dict(
             grant_type='authorization_code', client_id=self.client_id,
             redirect_uri=self.redirect_uri_oob, code=code)
@@ -83,13 +84,43 @@ class TestAuthFlow(unittest.TestCase):
 
         m.get(requests_mock.ANY, status_code=404)
         m.post(requests_mock.ANY, status_code=404)
-        m.post(self.token_endpoint, additional_matcher=post_body_matcher, json=token)
+        m.post(self.token_endpoint, additional_matcher=post_body_matcher, json=self.token)
 
         # verify token fetch flow
         flow = AuthFlow(**self.test_args)
 
         response = flow.fetch_token(code=code)
-        self.assertEqual(response, token)
+        self.assertEqual(response, self.token)
+
+        # verify token proxy to oauth2 session
+        self.assertEqual(flow.token, self.token)
+
+    def test_token_setter(self):
+        flow = AuthFlow(**self.test_args)
+
+        self.assertIsNone(flow.session.token)
+
+        flow.token = self.token
+
+        self.assertIsNotNone(flow.session.token)
+        self.assertEqual(flow.session.token['access_token'], self.token['access_token'])
+
+    def test_refresh_token(self):
+        flow = AuthFlow(**self.test_args)
+
+        with mock.patch.object(flow.session, 'refresh_token') as m:
+            flow.refresh_token()
+
+        m.assert_called_once_with(url=flow.token_endpoint)
+
+    def test_ensure_active_token(self):
+        flow = AuthFlow(**self.test_args)
+        flow.token = self.token
+
+        with mock.patch.object(flow.session, 'ensure_active_token') as m:
+            flow.ensure_active_token()
+
+        m.assert_called_once_with(token=self.token)
 
     def test_session_config(self):
         config = dict(
@@ -118,6 +149,16 @@ class TestAuthFlow(unittest.TestCase):
 
 class TestLeapAuthFlow(unittest.TestCase):
 
+    def test_from_default_config(self):
+        config = ClientConfig()
+
+        flow = LeapAuthFlow.from_config_model(config)
+
+        # endpoint urls are generated?
+        prefix = urljoin(DEFAULT_LEAP_API_ENDPOINT, '/')
+        self.assertTrue(flow.authorization_endpoint.startswith(prefix))
+        self.assertTrue(flow.token_endpoint.startswith(prefix))
+
     def test_from_minimal_config(self):
         config = ClientConfig(leap_api_endpoint='https://example.com/leap')
 
@@ -157,6 +198,15 @@ class TestLeapAuthFlow(unittest.TestCase):
         self.assertEqual(flow.session.headers.get('injected'), 'value')
         self.assertEqual(flow.session.default_timeout, 10)
 
+    def test_client_id_from_config(self):
+        client_id = '123'
+        config = ClientConfig(leap_api_endpoint='https://example.com/leap',
+                              leap_client_id=client_id)
+
+        flow = LeapAuthFlow.from_config_model(config)
+
+        self.assertEqual(flow.client_id, client_id)
+
 
 class TestLeapAuthFlowRunners(unittest.TestCase):
 
diff --git a/tests/test_regions.py b/tests/test_regions.py
index 731304d..b6af53d 100644
--- a/tests/test_regions.py
+++ b/tests/test_regions.py
@@ -13,15 +13,17 @@
 # limitations under the License.
 
 import unittest
+from secrets import token_hex
 from unittest import mock
 import contextlib
 
 from parameterized import parameterized
 
 from dwave.cloud.api.models import Region
+from dwave.cloud.api import constants
 from dwave.cloud.api.exceptions import ResourceAccessForbiddenError
 from dwave.cloud.config.models import ClientConfig
-from dwave.cloud.regions import get_regions
+from dwave.cloud.regions import get_regions, resolve_endpoints
 
 
 class GetRegionsInit(unittest.TestCase):
@@ -65,6 +67,9 @@ class GetRegionsFunctionality(unittest.TestCase):
             Region(code='A', name='Region A', endpoint='http://a/'),
             Region(code='B', name='Region B', endpoint='http://b/'),
         ]
+        # use a random mock metadata api to avoid cache contamination with mock data
+        # XXX: alternatively, we flush the cache after this test
+        mock_metadata_api_endpoint = f'https://example.com/meta/{token_hex(8)}/'
 
         # mock `api.Regions.from_config().list_regions()` call
         @contextlib.contextmanager
@@ -78,29 +83,29 @@ class GetRegionsFunctionality(unittest.TestCase):
                         side_effect=mock_regions_from_config) as m:
 
             with self.subTest('default call'):
-                ret = get_regions()
+                ret = get_regions(mock_metadata_api_endpoint)
                 m.assert_called()
                 self.assertEqual(ret, mock_regions)
 
             with self.subTest('caching'):
                 m.reset_mock()
-                ret = get_regions()
+                ret = get_regions(mock_metadata_api_endpoint)
                 m.assert_not_called()
                 self.assertEqual(ret, mock_regions)
 
             with self.subTest('cache refresh'):
                 m.reset_mock()
-                ret = get_regions(refresh=True)
+                ret = get_regions(mock_metadata_api_endpoint, refresh=True)
                 m.assert_called()
 
             with self.subTest('cache refresh'):
                 m.reset_mock()
-                ret = get_regions(maxage=0)
+                ret = get_regions(mock_metadata_api_endpoint, maxage=0)
                 m.assert_called()
 
             with self.subTest('cache bypass'):
                 m.reset_mock()
-                ret = get_regions(no_cache=True)
+                ret = get_regions(mock_metadata_api_endpoint, no_cache=True)
                 m.assert_called()
 
     @mock.patch("dwave.cloud.regions._fetch_available_regions",
@@ -108,3 +113,64 @@ class GetRegionsFunctionality(unittest.TestCase):
     def test_fetch_error(self, fetch_mock: mock.Mock):
         with self.assertRaises(ValueError):
             get_regions()
+
+
+class ResolveEndpointsMocked(unittest.TestCase):
+
+    def setUp(self):
+        self.mock_regions = [
+            Region(code='A', name='Region A', endpoint='https://a.example.com/path'),
+            Region(code='B', name='Region B', endpoint='https://b.example.com/path'),
+        ]
+
+        self.mocker = mock.patch('dwave.cloud.regions.get_regions',
+                                 return_value=self.mock_regions)
+        self.mocker.start()
+
+    def tearDown(self):
+        self.mocker.stop()
+
+    def test_null(self):
+        config = ClientConfig(metadata_api_endpoint=None, region=None,
+                              endpoint=None, leap_api_endpoint=None)
+
+        with mock.patch('dwave.cloud.regions.DEFAULT_REGION',
+                        self.mock_regions[0].code):
+            resolve_endpoints(config, inplace=True, shortcircuit=False)
+
+        self.assertEqual(config.metadata_api_endpoint, constants.DEFAULT_METADATA_API_ENDPOINT)
+        self.assertEqual(config.region, self.mock_regions[0].code)
+        self.assertEqual(config.endpoint, self.mock_regions[0].endpoint)
+        self.assertEqual(config.leap_api_endpoint, self.mock_regions[0].leap_api_endpoint)
+
+    def test_endpoint_override(self):
+        endpoint = 'https://example.com/sapi'
+        config = ClientConfig(endpoint=endpoint)
+
+        resolve_endpoints(config, inplace=True)
+
+        self.assertEqual(config.metadata_api_endpoint, constants.DEFAULT_METADATA_API_ENDPOINT)
+        self.assertIsNone(config.region)
+        self.assertEqual(config.endpoint, endpoint)
+        self.assertIsNotNone(config.leap_api_endpoint)
+
+    def test_leap_api_endpoint_override(self):
+        leap_api_endpoint = 'https://example.com/leap'
+        config = ClientConfig(leap_api_endpoint=leap_api_endpoint)
+
+        resolve_endpoints(config, inplace=True)
+
+        self.assertEqual(config.metadata_api_endpoint, constants.DEFAULT_METADATA_API_ENDPOINT)
+        self.assertIsNone(config.region)
+        self.assertEqual(config.leap_api_endpoint, leap_api_endpoint)
+        self.assertIsNotNone(config.endpoint)
+
+    def test_region_given(self):
+        mock_region = self.mock_regions[1].code
+        config = ClientConfig(region=mock_region)
+
+        resolve_endpoints(config, inplace=True)
+
+        self.assertEqual(config.region, mock_region)
+        self.assertEqual(config.endpoint, self.mock_regions[1].endpoint)
+        self.assertEqual(config.leap_api_endpoint, self.mock_regions[1].leap_api_endpoint)