Repository: dwavesystems/dwave-cloud-client

Issue:
Add support for custom headers
It's useful during development and testing to be able to pass-in a custom HTTP header in every API call.

Add support for `headers` config parameter (file, constructor) and `DWAVE_API_HEADERS` environment variable.

Comments:

Patch:
diff --git a/dwave/cloud/client.py b/dwave/cloud/client.py
index b13cb8e..4055220 100644
--- a/dwave/cloud/client.py
+++ b/dwave/cloud/client.py
@@ -114,6 +114,9 @@ class Client(object):
         connection_close (bool, default=False):
             Force HTTP(S) connection close after each request.
 
+        headers (dict/str):
+            Additional HTTP headers.
+
     Other Parameters:
         Unrecognized keys (str):
             All unrecognized keys are passed through to the appropriate client class constructor
@@ -183,7 +186,7 @@ class Client(object):
     @classmethod
     def from_config(cls, config_file=None, profile=None, client=None,
                     endpoint=None, token=None, solver=None, proxy=None,
-                    legacy_config_fallback=False, **kwargs):
+                    headers=None, legacy_config_fallback=False, **kwargs):
         """Client factory method to instantiate a client instance from configuration.
 
         Configuration values can be specified in multiple ways, ranked in the following
@@ -211,8 +214,9 @@ class Client(object):
         file is not explicitly specified, detected on the system, or defined via
         an environment variable.
 
-        Environment variables: ``DWAVE_CONFIG_FILE``, ``DWAVE_PROFILE``, ``DWAVE_API_CLIENT``,
-        ``DWAVE_API_ENDPOINT``, ``DWAVE_API_TOKEN``, ``DWAVE_API_SOLVER``, ``DWAVE_API_PROXY``.
+        Environment variables: ``DWAVE_CONFIG_FILE``, ``DWAVE_PROFILE``,
+        ``DWAVE_API_CLIENT``, ``DWAVE_API_ENDPOINT``, ``DWAVE_API_TOKEN``,
+        ``DWAVE_API_SOLVER``, ``DWAVE_API_PROXY``, ``DWAVE_API_HEADERS``.
 
         Environment variables are described in :mod:`dwave.cloud.config`.
 
@@ -268,6 +272,10 @@ class Client(object):
                 username/password, port, scheme, etc. If undefined, client
                 uses the system-level proxy, if defined, or connects directly to the API.
 
+            headers (dict/str, default=None):
+                Newline-separated additional HTTP headers to include with each
+                API request, or a dictionary of (key, value) pairs.
+
             legacy_config_fallback (bool, default=False):
                 If True and loading from a standard D-Wave Cloud Client configuration
                 file (``dwave.conf``) fails, tries loading a legacy configuration file (``~/.dwrc``).
@@ -292,7 +300,6 @@ class Client(object):
                 Config file parse failed.
 
         Examples:
-
             A variety of examples are given in :mod:`dwave.cloud.config`.
 
             This example initializes :class:`~dwave.cloud.client.Client` from an
@@ -309,7 +316,8 @@ class Client(object):
         # (`./dwave.conf`, `~/.config/dwave/dwave.conf`, etc)
         config = load_config(
             config_file=config_file, profile=profile, client=client,
-            endpoint=endpoint, token=token, solver=solver, proxy=proxy)
+            endpoint=endpoint, token=token, solver=solver, proxy=proxy,
+            headers=headers)
         logger.debug("Config loaded: %r", config)
 
         # fallback to legacy `.dwrc` if key variables missing
@@ -319,8 +327,8 @@ class Client(object):
 
             if not config.get('token'):
                 config = legacy_load_config(
-                    profile=profile, client=client,
-                    endpoint=endpoint, token=token, solver=solver, proxy=proxy)
+                    profile=profile, client=client, endpoint=endpoint,
+                    token=token, solver=solver, proxy=proxy, headers=headers)
                 logger.debug("Legacy config loaded: %r", config)
 
         # manual override of other (client-custom) arguments
@@ -335,7 +343,7 @@ class Client(object):
 
     def __init__(self, endpoint=None, token=None, solver=None, proxy=None,
                  permissive_ssl=False, request_timeout=60, polling_timeout=None,
-                 connection_close=False, **kwargs):
+                 connection_close=False, headers=None, **kwargs):
         """To setup the connection a pipeline of queues/workers is constructed.
 
         There are five interactions with the server the connection manages:
@@ -357,16 +365,18 @@ class Client(object):
 
         logger.debug(
             "Creating a client for (endpoint=%r, token=%r, solver=%r, proxy=%r, "
-            "permissive_ssl=%r, request_timeout=%r, polling_timeout=%r, **kwargs=%r)",
-            endpoint, token, solver, proxy, permissive_ssl, request_timeout, polling_timeout, kwargs
+            "permissive_ssl=%r, request_timeout=%r, polling_timeout=%r, "
+            "connection_close=%r, headers=%r, **kwargs=%r)",
+            endpoint, token, solver, proxy,
+            permissive_ssl, request_timeout, polling_timeout,
+            connection_close, headers, kwargs
         )
 
+        # parse solver
         if not solver:
             solver_def = {}
-
         elif isinstance(solver, collections.Mapping):
             solver_def = solver
-
         elif isinstance(solver, six.string_types):
             # support features dict encoded as JSON in our config INI file
             # TODO: push this decoding to the config module, once we switch to a
@@ -379,9 +389,27 @@ class Client(object):
                 # features dict (equality constraint on full solver name)
                 logger.debug("Invalid solver JSON, assuming string name: %r", solver)
                 solver_def = dict(name__eq=solver)
-
         else:
             raise ValueError("Expecting a features dictionary or a string name for 'solver'")
+        logger.debug("Parsed solver=%r", solver_def)
+
+        # parse headers
+        if not headers:
+            headers_dict = {}
+        elif isinstance(headers, collections.Mapping):
+            headers_dict = headers
+        elif isinstance(headers, six.string_types):
+            try:
+                # valid  headers = "Field-1: value-1\nField-2: value-2"
+                headers_dict = {key.strip(): val.strip()
+                                for key, val in [line.split(':')
+                                                 for line in headers.strip().split('\n')]}
+            except Exception as e:
+                logger.debug("Invalid headers: %r", headers)
+                headers_dict = {}
+        else:
+            raise ValueError("HTTP headers expected in a dict, or a string")
+        logger.debug("Parsed headers=%r", headers_dict)
 
         # Store connection/session parameters
         self.endpoint = endpoint
@@ -392,6 +420,7 @@ class Client(object):
         self.polling_timeout = parse_float(polling_timeout)
 
         self.proxy = proxy
+        self.headers = headers_dict
         self.permissive_ssl = permissive_ssl
         self.connection_close = connection_close
 
@@ -456,6 +485,8 @@ class Client(object):
         session = BaseUrlSession(base_url=endpoint)
         session.mount('http://', TimeoutingHTTPAdapter(timeout=self.request_timeout))
         session.mount('https://', TimeoutingHTTPAdapter(timeout=self.request_timeout))
+        if self.headers:
+            session.headers.update(self.headers)
         session.headers.update({'X-Auth-Token': self.token,
                                 'User-Agent': user_agent(__packagename__, __version__)})
         session.proxies = {'http': self.proxy, 'https': self.proxy}
diff --git a/dwave/cloud/config.py b/dwave/cloud/config.py
index cf26496..7bdaf70 100644
--- a/dwave/cloud/config.py
+++ b/dwave/cloud/config.py
@@ -59,6 +59,8 @@ Environment variables:
 
     ``DWAVE_API_PROXY``: URL for proxy connections to D-Wave API.
 
+    ``DWAVE_API_HEADERS``: Optional additional HTTP headers.
+
 Examples:
     The following are typical examples of using :func:`~dwave.cloud.client.Client.from_config`
     to create a configured client.
@@ -195,12 +197,14 @@ Examples:
     {'client': u'sw',
      'endpoint': u'https://url.of.some.dwavesystem.com/sapi',
      'proxy': None,
+     'headers': None,
      'solver': 'EXAMPLE_2000Q_SYSTEM',
      'token': u'ABC-123456789123456789123456789'}
     >>> dc.config.load_config("./dwave_c.conf", profile='dw2000b', solver='Solver3')   # doctest: +SKIP
     {'client': u'qpu',
      'endpoint': u'https://url.of.some.dwavesystem.com/sapi',
      'proxy': None,
+     'headers': None,
      'solver': 'Solver3',
      'token': u'DEF-987654321987654321987654321'}
 
@@ -540,18 +544,18 @@ def load_profile_from_files(filenames=None, profile=None):
         The following example code loads profile values from parsing both these files,
         by default loading the first profile encountered or an explicitly specified profile.
 
-        >>> import dwave.cloud as dc
-        >>> dc.config.load_profile_from_files(["./dwave_a.conf", "./dwave_b.conf"])   # doctest: +SKIP
-        {'client': u'sw',
-         'endpoint': u'https://url.of.some.dwavesystem.com/sapi',
-         'solver': u'EXAMPLE_2000Q_SYSTEM_A',
-         'token': u'DEF-987654321987654321987654321'}
-        >>> dc.config.load_profile_from_files(["./dwave_a.conf", "./dwave_b.conf"],
-        ...                                   profile='dw2000b')   # doctest: +SKIP
-        {'client': u'qpu',
-        'endpoint': u'https://url.of.some.other.dwavesystem.com/sapi',
-        'solver': u'EXAMPLE_2000Q_SYSTEM_B',
-        'token': u'ABC-123456789123456789123456789'}
+        >>> from dwave.cloud import config
+        >>> config.load_profile_from_files(["./dwave_a.conf", "./dwave_b.conf"])   # doctest: +SKIP
+        {'client': 'sw',
+         'endpoint': 'https://url.of.some.dwavesystem.com/sapi',
+         'solver': 'EXAMPLE_2000Q_SYSTEM_A',
+         'token': 'DEF-987654321987654321987654321'}
+        >>> config.load_profile_from_files(["./dwave_a.conf", "./dwave_b.conf"],
+        ...                                profile='dw2000b')   # doctest: +SKIP
+        {'client': 'qpu',
+        'endpoint': 'https://url.of.some.other.dwavesystem.com/sapi',
+        'solver': 'EXAMPLE_2000Q_SYSTEM_B',
+        'token': 'ABC-123456789123456789123456789'}
 
     """
 
@@ -617,7 +621,8 @@ def get_default_config():
 
 
 def load_config(config_file=None, profile=None, client=None,
-                endpoint=None, token=None, solver=None, proxy=None):
+                endpoint=None, token=None, solver=None,
+                proxy=None, headers=None):
     """Load D-Wave Cloud Client configuration based on a configuration file.
 
     Configuration values can be specified in multiple ways, ranked in the following
@@ -648,13 +653,13 @@ def load_config(config_file=None, profile=None, client=None,
     file is not explicitly specified, detected on the system, or defined via
     an environment variable.
 
-    Environment variables: ``DWAVE_CONFIG_FILE``, ``DWAVE_PROFILE``, ``DWAVE_API_CLIENT``,
-    ``DWAVE_API_ENDPOINT``, ``DWAVE_API_TOKEN``, ``DWAVE_API_SOLVER``, ``DWAVE_API_PROXY``.
+    Environment variables: ``DWAVE_CONFIG_FILE``, ``DWAVE_PROFILE``,
+    ``DWAVE_API_CLIENT``, ``DWAVE_API_ENDPOINT``, ``DWAVE_API_TOKEN``,
+    ``DWAVE_API_SOLVER``, ``DWAVE_API_PROXY``, ``DWAVE_API_HEADERS``.
 
     Environment variables are described in :mod:`dwave.cloud.config`.
 
     Args:
-
         config_file (str/[str]/None/False/True, default=None):
             Path to configuration file(s).
 
@@ -689,7 +694,7 @@ def load_config(config_file=None, profile=None, client=None,
         token (str, default=None):
             API authorization token.
 
-        solver (str, default=None):
+        solver (dict/str, default=None):
             :term:`solver` features, as a JSON-encoded dictionary of feature constraints,
             the client should use. See :meth:`~dwave.cloud.client.Client.get_solvers` for
             semantics of supported feature constraints.
@@ -705,13 +710,16 @@ def load_config(config_file=None, profile=None, client=None,
             username/password, port, scheme, etc. If undefined, client
             uses the system-level proxy, if defined, or connects directly to the API.
 
+        headers (dict/str, default=None):
+            Header lines to include in API calls, each line formatted as
+             ``Key: value``, or a parsed dictionary.
+
     Returns:
         dict:
-            Mapping of configuration keys to values for the profile
-            (section), as read from the configuration file and optionally overridden by
-            environment values and specified keyword arguments.
-            Always contains the `client`, `endpoint`, `token`, `solver`, and `proxy`
-            keys.
+            Mapping of configuration keys to values for the profile (section),
+            as read from the configuration file and optionally overridden by
+            environment values and specified keyword arguments. Always contains
+            the `client`, `endpoint`, `token`, `solver`, and `proxy` keys.
 
     Raises:
         :exc:`ValueError`:
@@ -727,16 +735,17 @@ def load_config(config_file=None, profile=None, client=None,
         This example loads the configuration from an auto-detected configuration file
         in the home directory of a Windows system user.
 
-        >>> import dwave.cloud as dc
-        >>> dc.config.load_config()
-        {'client': u'qpu',
-         'endpoint': u'https://url.of.some.dwavesystem.com/sapi',
+        >>> from dwave.cloud import config
+        >>> config.load_config()
+        {'client': 'qpu',
+         'endpoint': 'https://url.of.some.dwavesystem.com/sapi',
          'proxy': None,
-         'solver': u'EXAMPLE_2000Q_SYSTEM_A',
-         'token': u'DEF-987654321987654321987654321'}
+         'solver': 'EXAMPLE_2000Q_SYSTEM_A',
+         'token': 'DEF-987654321987654321987654321',
+         'headers': None}
         >>> See which configuration file was loaded
-        >>> dc.config.get_configfile_paths()
-        [u'C:\\Users\\jane\\AppData\\Local\\dwavesystem\\dwave\\dwave.conf']
+        >>> config.get_configfile_paths()
+        ['C:\\Users\\jane\\AppData\\Local\\dwavesystem\\dwave\\dwave.conf']
 
         Additional examples are given in :mod:`dwave.cloud.config`.
 
@@ -774,6 +783,7 @@ def load_config(config_file=None, profile=None, client=None,
     section['token'] = token or os.getenv("DWAVE_API_TOKEN", section.get('token'))
     section['solver'] = solver or os.getenv("DWAVE_API_SOLVER", section.get('solver'))
     section['proxy'] = proxy or os.getenv("DWAVE_API_PROXY", section.get('proxy'))
+    section['headers'] = headers or os.getenv("DWAVE_API_HEADERS", section.get('headers'))
 
     return section
 


Test patch:
diff --git a/tests/test_client.py b/tests/test_client.py
index 009e167..1b5145c 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -230,6 +230,35 @@ class ClientFactory(unittest.TestCase):
             with dwave.cloud.Client.from_config(solver='solver') as client:
                 self.assertEqual(client.default_solver, {"name__eq": "solver"})
 
+    def test_headers_from_config(self):
+        headers_dict = {"key-1": "value-1", "key-2": "value-2"}
+        headers_str = """  key-1:value-1
+            key-2: value-2
+        """
+        conf = dict(token='token', headers=headers_str)
+
+        with mock.patch("dwave.cloud.client.load_config", lambda **kwargs: conf):
+            with dwave.cloud.Client.from_config() as client:
+                self.assertDictEqual(client.headers, headers_dict)
+
+    def test_headers_from_kwargs(self):
+        headers_dict = {"key-1": "value-1", "key-2": "value-2"}
+        headers_str = "key-2:value-2\nkey-1:value-1"
+        conf = dict(token='token')
+
+        def load_config(**kwargs):
+            return merge(kwargs, conf, op=lambda a, b: a or b)
+
+        # headers as dict
+        with mock.patch("dwave.cloud.client.load_config", load_config):
+            with dwave.cloud.Client.from_config(headers=headers_dict) as client:
+                self.assertDictEqual(client.headers, headers_dict)
+
+        # headers as str
+        with mock.patch("dwave.cloud.client.load_config", load_config):
+            with dwave.cloud.Client.from_config(headers=headers_str) as client:
+                self.assertDictEqual(client.headers, headers_dict)
+
 
 class FeatureBasedSolverSelection(unittest.TestCase):
     """Test Client.get_solvers(**filters)."""
diff --git a/tests/test_config.py b/tests/test_config.py
index 20e4914..2b34bea 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -47,6 +47,8 @@ class TestConfig(unittest.TestCase):
         endpoint = https://url.to.alpha/api
         proxy = http://user:pass@myproxy.com:8080/
         token = alpha-token
+        headers =  key-1:value-1
+          key-2: value-2
     """
 
     def parse_config_string(self, text):
@@ -127,6 +129,8 @@ class TestConfig(unittest.TestCase):
         self.assertEqual(config['endpoint'], "https://url.to.alpha/api")
         # default values are inherited
         self.assertEqual(config['client'], "qpu")
+        # multi-line values are read
+        self.assertEqual(config['headers'], "key-1:value-1\nkey-2: value-2")
 
     def _load_config_from_files(self, asked, provided=None, data=None):
         self.assertEqual(asked, provided)
@@ -210,6 +214,7 @@ class TestConfig(unittest.TestCase):
                 self.assertEqual(load_config(client='manual')['client'], 'manual')
                 self.assertEqual(load_config(solver='manual')['solver'], 'manual')
                 self.assertEqual(load_config(proxy='manual')['proxy'], 'manual')
+                self.assertEqual(load_config(headers='headers')['headers'], 'headers')
 
     def test_config_load__profile_arg_nonexisting(self):
         """load_config should fail if the profile specified in kwargs or env in
@@ -332,3 +337,25 @@ class TestConfig(unittest.TestCase):
             m.assert_has_calls([mock.call('file1', 'r'), mock.call('file2', 'r')])
             self.assertEqual(section['endpoint'], 'alpha')
             self.assertEqual(section['solver'], 'DW_2000Q_2')
+
+    def test_config_load_env_override(self):
+        with mock.patch("dwave.cloud.config.load_config_from_files",
+                        partial(self._load_config_from_files, data=u"", provided=['myfile'])):
+
+            with mock.patch.dict(os.environ, {'DWAVE_API_CLIENT': 'test'}):
+                self.assertEqual(load_config(config_file='myfile')['client'], 'test')
+
+            with mock.patch.dict(os.environ, {'DWAVE_API_ENDPOINT': 'test'}):
+                self.assertEqual(load_config(config_file='myfile')['endpoint'], 'test')
+
+            with mock.patch.dict(os.environ, {'DWAVE_API_TOKEN': 'test'}):
+                self.assertEqual(load_config(config_file='myfile')['token'], 'test')
+
+            with mock.patch.dict(os.environ, {'DWAVE_API_SOLVER': 'test'}):
+                self.assertEqual(load_config(config_file='myfile')['solver'], 'test')
+
+            with mock.patch.dict(os.environ, {'DWAVE_API_PROXY': 'test'}):
+                self.assertEqual(load_config(config_file='myfile')['proxy'], 'test')
+
+            with mock.patch.dict(os.environ, {'DWAVE_API_HEADERS': 'test'}):
+                self.assertEqual(load_config(config_file='myfile')['headers'], 'test')