Repository: dwavesystems/dwave-cloud-client

Issue:
Add long polling support
Recently SAPI added support for long polling calls on problem status endpoints.

Long polling behavior is activated by appending `timeout=<sec>` query parameter.

Default value for timeout is 1 second, and maximum allowed value is 30 seconds.

In addition to short polling with exponential back-off, we want to add long-polling as polling schedule strategy, and eventually make it the default.

Discussion:

Patch:
diff --git a/dwave/cloud/api/resources.py b/dwave/cloud/api/resources.py
index d18c5cd..a0487d7 100644
--- a/dwave/cloud/api/resources.py
+++ b/dwave/cloud/api/resources.py
@@ -185,17 +185,25 @@ class Problems(ResourceBase):
         return models.ProblemStatusMaybeWithAnswer.model_validate(status)
 
     @accepts(media_type='application/vnd.dwave.sapi.problems+json', version='>=2.1,<3')
-    def get_problem_status(self, problem_id: str) -> models.ProblemStatus:
+    def get_problem_status(self,
+                           problem_id: str,
+                           timeout: Optional[int] = None,
+                           ) -> models.ProblemStatus:
         """Retrieve short status of a single problem."""
         path = ''
         params = dict(id=problem_id)
+        if timeout is not None:
+            params.setdefault('timeout', timeout)
         response = self.session.get(path, params=params)
         status = response.json()[0]
         return models.ProblemStatus.model_validate(status)
 
     # XXX: @pydantic.validate_arguments
     @accepts(media_type='application/vnd.dwave.sapi.problems+json', version='>=2.1,<3')
-    def get_problem_statuses(self, problem_ids: List[str]) -> List[models.ProblemStatus]:
+    def get_problem_statuses(self,
+                             problem_ids: List[str],
+                             timeout: Optional[int] = None,
+                             ) -> List[models.ProblemStatus]:
         """Retrieve short problem statuses for a list of problems."""
         if not isinstance(problem_ids, list):
             raise TypeError('a list of problem ids expected')
@@ -204,6 +212,8 @@ class Problems(ResourceBase):
 
         path = ''
         params = dict(id=','.join(problem_ids))
+        if timeout is not None:
+            params.setdefault('timeout', timeout)
         response = self.session.get(path, params=params)
         statuses = response.json()
         return TypeAdapter(List[models.ProblemStatus]).validate_python(statuses)
diff --git a/dwave/cloud/client/base.py b/dwave/cloud/client/base.py
index 4bc4ef6..b455c77 100644
--- a/dwave/cloud/client/base.py
+++ b/dwave/cloud/client/base.py
@@ -66,7 +66,7 @@ from dwave.cloud.package_info import __packagename__, __version__
 from dwave.cloud.exceptions import *    # TODO: fix
 from dwave.cloud.computation import Future
 from dwave.cloud.config import load_config, update_config, validate_config_v1
-from dwave.cloud.config.models import ClientConfig
+from dwave.cloud.config.models import ClientConfig, PollingStrategy
 from dwave.cloud.solver import available_solvers, StructuredSolver, UnstructuredSolver
 from dwave.cloud.concurrency import PriorityThreadPoolExecutor
 from dwave.cloud.regions import get_regions, resolve_endpoints
@@ -307,10 +307,14 @@ class Client(object):
         'headers': None,
         'client_cert': None,
         'client_cert_key': None,
+        'poll_strategy': 'backoff',
         # poll back-off schedule defaults [sec]
         'poll_backoff_min': 0.05,
         'poll_backoff_max': 60,
         'poll_backoff_base': 1.3,
+        # long polling parameters
+        'poll_wait_time': 30,
+        'poll_pause': 0,
         # idempotent http requests retry params
         'http_retry_total': 10,
         'http_retry_connect': None,
@@ -334,7 +338,7 @@ class Client(object):
     _UPLOAD_PART_THREAD_COUNT = 10
     _ENCODE_PROBLEM_THREAD_COUNT = _UPLOAD_PROBLEM_THREAD_COUNT
     _CANCEL_THREAD_COUNT = 1
-    _POLL_THREAD_COUNT = 2
+    _POLL_THREAD_COUNT = 5
     _LOAD_THREAD_COUNT = 5
 
     # Poll grouping time frame; two scheduled polls are grouped if closer than [sec]:
@@ -1449,9 +1453,18 @@ class Client(object):
         finally:
             session.close()
 
-    def _poll(self, future):
+    def _poll(self, future: Future) -> None:
         """Enqueue a problem to poll the server for status."""
 
+        # split for simplicity
+        if self.config.polling_schedule.strategy == PollingStrategy.BACKOFF:
+            return self._poll_using_backoff(future)
+        elif self.config.polling_schedule.strategy == PollingStrategy.LONG_POLLING:
+            return self._poll_using_long_polling(future)
+        else:
+            raise RuntimeError("unexpected polling strategy")
+
+    def _poll_using_backoff(self, future: Future) -> None:
         if future._poll_backoff is None:
             # on first poll, start with minimal back-off
             future._poll_backoff = self.config.polling_schedule.backoff_min
@@ -1479,6 +1492,13 @@ class Client(object):
 
         self._poll_queue.put((at, future))
 
+    def _poll_using_long_polling(self, future: Future) -> None:
+        # we use problem submit time to prioritize polling of jobs submitted earlier
+        at = datetime_to_timestamp(future.time_created)
+
+        # use the same priority queue as for backoff polling
+        self._poll_queue.put((at, future))
+
     def _do_poll_problems(self):
         """Poll the server for the status of a set of problems.
 
@@ -1488,8 +1508,12 @@ class Client(object):
         session = self.create_session()
         try:
             # grouped futures (all scheduled within _POLL_GROUP_TIMEFRAME)
+            # and/or up to _STATUS_QUERY_SIZE (depending on strategy)
             frame_futures = {}
 
+            use_long_polling = (
+                self.config.polling_schedule.strategy == PollingStrategy.LONG_POLLING)
+
             def task_done():
                 self._poll_queue.task_done()
 
@@ -1518,6 +1542,7 @@ class Client(object):
                     return
 
                 # try grouping if scheduled within grouping timeframe
+                # (or in the long polling case, add up to _STATUS_QUERY_SIZE available futures)
                 while len(frame_futures) < self._STATUS_QUERY_SIZE:
                     try:
                         task = self._poll_queue.get_nowait()
@@ -1525,7 +1550,7 @@ class Client(object):
                         break
 
                     at, future = task
-                    if at - frame_earliest <= self._POLL_GROUP_TIMEFRAME:
+                    if use_long_polling or (at - frame_earliest <= self._POLL_GROUP_TIMEFRAME):
                         if not add(future):
                             return
                     else:
@@ -1537,18 +1562,21 @@ class Client(object):
                 ids = [future.id for future in frame_futures.values()]
                 logger.debug("Polling for status of futures: %s", ids)
                 query_string = 'problems/?id=' + ','.join(ids)
+                if use_long_polling:
+                    query_string += f'&timeout={self.config.polling_schedule.wait_time}'
 
                 # if futures were cancelled while `add`ing, skip empty frame
                 if not ids:
                     continue
 
-                # wait until `frame_earliest` before polling
-                delay = frame_earliest - time.time()
-                if delay > 0:
-                    logger.debug("Pausing polling %.2f sec for futures: %s", delay, ids)
-                    time.sleep(delay)
-                else:
-                    logger.trace("Skipping non-positive delay of %.2f sec", delay)
+                if not use_long_polling:
+                    # wait until `frame_earliest` before polling
+                    delay = frame_earliest - time.time()
+                    if delay > 0:
+                        logger.debug("Pausing polling %.2f sec for futures: %s", delay, ids)
+                        time.sleep(delay)
+                    else:
+                        logger.trace("Skipping non-positive delay of %.2f sec", delay)
 
                 # execute and handle the polling request
                 try:
@@ -1584,6 +1612,10 @@ class Client(object):
                 for id_ in frame_futures.keys():
                     task_done()
 
+                if use_long_polling and (pause := self.config.polling_schedule.pause) > 0:
+                    logger.debug("Pausing %.3f sec between long polling requests", pause)
+                    time.sleep(pause)
+
         except Exception as err:
             logger.exception(err)
 
diff --git a/dwave/cloud/config/models.py b/dwave/cloud/config/models.py
index eed987f..4bec0f1 100644
--- a/dwave/cloud/config/models.py
+++ b/dwave/cloud/config/models.py
@@ -14,6 +14,7 @@
 
 import ast
 import copy
+import enum
 import json
 import logging
 from collections import abc
@@ -21,13 +22,14 @@ from typing import Optional, Union, Literal, Tuple, Dict, Any
 from typing_extensions import Annotated     # backport for py37, py38
 
 import urllib3
-from pydantic import BaseModel, BeforeValidator
+from pydantic import BaseModel, BeforeValidator, NonNegativeInt
 
 from dwave.cloud.config.loaders import update_config
 from dwave.cloud.api.constants import (
     DEFAULT_REGION, DEFAULT_METADATA_API_ENDPOINT, DEFAULT_LEAP_API_ENDPOINT)
 
-__all__ = ['RequestRetryConfig', 'PollingSchedule', 'ClientConfig',
+__all__ = ['RequestRetryConfig', 'ClientConfig',
+           'BackoffPollingSchedule', 'LongPollingSchedule',
            'validate_config_v1', 'dump_config_v1', 'load_config_v1']
 
 logger = logging.getLogger(__name__)
@@ -87,9 +89,17 @@ class RequestRetryConfig(BaseModel, GetterMixin):
         return retry
 
 
-class PollingSchedule(BaseModel, GetterMixin):
+class PollingStrategy(str, enum.Enum):
+    BACKOFF = "backoff"
+    LONG_POLLING = "long-polling"
+
+
+class BackoffPollingSchedule(BaseModel):
     """Problem status polling exponential back-off schedule params."""
 
+    #: Polling with exponential backoff
+    strategy: Literal[PollingStrategy.BACKOFF] = PollingStrategy.BACKOFF
+
     #: Duration of the first interval (between first and second poll), in seconds.
     backoff_min: Optional[float] = 0.05
 
@@ -101,6 +111,20 @@ class PollingSchedule(BaseModel, GetterMixin):
     backoff_base: Optional[float] = 1.3
 
 
+class LongPollingSchedule(BaseModel):
+    """Problem status long polling params."""
+
+    #: Long polling strategy
+    strategy: Literal[PollingStrategy.LONG_POLLING] = PollingStrategy.LONG_POLLING
+
+    #: Maximum duration a long polling connection is kept open, in whole number of seconds.
+    #: Note: SAPI requires a non-negative integer wait_time
+    wait_time: Optional[NonNegativeInt] = 30
+
+    #: Pause between two successive long polling connections, in seconds.
+    pause: Optional[float] = 0.0
+
+
 def _literal_eval(obj):
     if isinstance(obj, str):
         return ast.literal_eval(obj)
@@ -124,8 +148,11 @@ class ClientConfig(BaseModel, GetterMixin):
     client: Optional[str] = None
     solver: Optional[Dict[str, Any]] = None
 
-    # [sapi client specific] poll back-off schedule defaults [sec]
-    polling_schedule: Optional[PollingSchedule] = PollingSchedule()
+    # [sapi client specific] polling schedule defaults [sec]
+    # note: discriminated unions are faster than unions, but we want to be able
+    # to set a defautl value when `strategy` is not specified
+    polling_schedule: Optional[Union[BackoffPollingSchedule,
+                                     LongPollingSchedule]] = BackoffPollingSchedule()
     polling_timeout: Optional[float] = None
 
     # general http(s) connection params
@@ -220,6 +247,9 @@ def validate_config_v1(raw_config: dict) -> ClientConfig:
     config['polling_schedule'] = {k[len(prefix):]: v for k, v in config.items()
                                   if k.startswith(prefix)}
 
+    # set a default polling strategy
+    config['polling_schedule'].setdefault('strategy', PollingStrategy.BACKOFF.value)
+
     # retry
     prefix = 'http_retry_'
     config['request_retry'] = {k[len(prefix):]: v for k, v in config.items()
@@ -285,10 +315,14 @@ _V1_CONFIG_DEFAULTS = {
     'headers': None,
     'client_cert': None,
     'client_cert_key': None,
+    'poll_strategy': 'backoff',
     # poll back-off schedule defaults
     'poll_backoff_min': 0.05,
     'poll_backoff_max': 60,
     'poll_backoff_base': 1.3,
+    # long polling parameters
+    'poll_wait_time': 30,
+    'poll_pause': 0,
     # idempotent http requests retry params
     'http_retry_total': 10,
     'http_retry_connect': None,
diff --git a/releasenotes/notes/add-long-polling-support-943bd017d05e563c.yaml b/releasenotes/notes/add-long-polling-support-943bd017d05e563c.yaml
new file mode 100644
index 0000000..4f9a34d
--- /dev/null
+++ b/releasenotes/notes/add-long-polling-support-943bd017d05e563c.yaml
@@ -0,0 +1,16 @@
+---
+features:
+  - |
+    Add long polling as problem status/answer polling strategy.
+
+    We add the following config options: ``poll_strategy`` (``"backoff"`` is the
+    current default, ``"long-polling"`` is the new strategy), ``poll_wait_time``,
+    and ``poll_pause``.
+
+  - |
+    Add ``timeout`` keyword argument to ``api.resources.Problems.get_problem_status()``
+    and ``api.resources.Problems.get_problem_statuses()`` methods.
+
+    Use it to enable long polling on problem status API endpoint. The connection
+    is closed (server-side) when at least one problem completes, or the ``timeout``
+    (in seconds) elapses, whatever comes first.


Test patch:
diff --git a/tests/api/resources/test_problems.py b/tests/api/resources/test_problems.py
index c22b2a8..26cf73a 100644
--- a/tests/api/resources/test_problems.py
+++ b/tests/api/resources/test_problems.py
@@ -160,6 +160,15 @@ class ProblemResourcesBaseTests(abc.ABC):
         self.assertIsInstance(status, models.ProblemStatus)
         self.verify_problem_status(status, solved=True)
 
+    def test_get_problem_status_long_polling(self):
+        """Problem status retrieved by problem id, long polling variant."""
+
+        problem_id = self.problem_id
+        status = self.api.get_problem_status(problem_id, timeout=self.poll_wait_time)
+
+        self.assertIsInstance(status, models.ProblemStatus)
+        self.verify_problem_status(status, solved=True)
+
     def test_get_problem_statuses(self):
         """Multiple problem statuses retrieved by problem ids."""
 
@@ -172,6 +181,18 @@ class ProblemResourcesBaseTests(abc.ABC):
             self.assertIsInstance(status, models.ProblemStatus)
             self.verify_problem_status(status, solved=True)
 
+    def test_get_problem_statuses_long_polling(self):
+        """Multiple problem statuses retrieved by problem ids, long polling variant."""
+
+        problem_id = self.problem_id
+        ids = [problem_id] * 3
+        statuses = self.api.get_problem_statuses(ids, timeout=self.poll_wait_time)
+
+        self.assertEqual(len(statuses), len(ids))
+        for status in statuses:
+            self.assertIsInstance(status, models.ProblemStatus)
+            self.verify_problem_status(status, solved=True)
+
     def test_get_problem_info(self):
         """Problem info (complete problem description) retrieved by problem id."""
 
@@ -427,6 +448,9 @@ class ProblemResourcesMockerMixin:
         all_solvers = [self.solver_id]
         all_statuses = {s.value for s in constants.ProblemStatus}
 
+        # long polling param
+        self.poll_wait_time = 20
+
         # problem status et al
         self.mocker.get(
             url(f'problems/{p1_id}'),
@@ -437,6 +461,11 @@ class ProblemResourcesMockerMixin:
             complete_qs=True,
             json=[p1_status],
             request_headers=headers)
+        self.mocker.get(
+            url(f'problems/?id={p1_id}&timeout={self.poll_wait_time}'),
+            complete_qs=True,
+            json=[p1_status],
+            request_headers=headers)
         self.mocker.get(
             url(f'problems/?id={p2_id}'),
             complete_qs=True,
@@ -468,6 +497,11 @@ class ProblemResourcesMockerMixin:
             complete_qs=True,
             json=[p1_status] * 3,
             request_headers=headers)
+        self.mocker.get(
+            url(f'problems/?id={p1_id},{p1_id},{p1_id}&timeout={self.poll_wait_time}'),
+            complete_qs=True,
+            json=[p1_status] * 3,
+            request_headers=headers)
 
         # list problems
         self.mocker.get(
@@ -653,7 +687,7 @@ class TestMockProblemsUnstructured(UnstructuredProblemTestsMixin,
 
 
 @unittest.skipUnless(dimod, "dimod not installed")
-class TestMockProblemsUnstructuredWithBinaryRefAsnwer(TestMockProblemsUnstructured):
+class TestMockProblemsUnstructuredWithBinaryRefAnswer(TestMockProblemsUnstructured):
     """Verify `dwave.cloud.api.resources.Problems` handle unstructured problems
     with answers in the `binary-ref` format, using mocked SAPI responses.
     """
@@ -691,6 +725,7 @@ class TestCloudProblemsStructured(StructuredProblemTestsMixin,
         with Client(**config) as client:
             cls.token = client.config.token
             cls.api = Problems.from_config(client.config)
+            cls.poll_wait_time = 2
 
             # submit and solve an Ising problem as a fixture
             solver = client.get_solver(qpu=True)
@@ -731,6 +766,7 @@ class TestCloudProblemsUnstructured(UnstructuredProblemTestsMixin,
         with Client(**config) as client:
             cls.token = client.config.token
             cls.api = Problems.from_config(client.config)
+            cls.poll_wait_time = 2
 
             # submit and solve an Ising problem as a fixture
             solver = client.get_solver(hybrid=True)
diff --git a/tests/test_client.py b/tests/test_client.py
index 94c2533..d3ad830 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -32,7 +32,7 @@ from plucky import merge
 
 from dwave.cloud.api import constants, models, Regions
 from dwave.cloud.client import Client
-from dwave.cloud.config.models import dump_config_v1
+from dwave.cloud.config.models import dump_config_v1, PollingStrategy
 from dwave.cloud.exceptions import (
     SolverAuthenticationError, SolverError, SolverNotFoundError)
 from dwave.cloud.solver import StructuredSolver, UnstructuredSolver
@@ -542,30 +542,53 @@ class ClientConstruction(unittest.TestCase):
                 self.assertEqual(client.config.cert, client_cert)
 
     def test_polling_params_from_config(self):
-        poll_conf = {"poll_backoff_min": "0.1", "poll_backoff_max": "1"}
-        conf = dict(token='token', **poll_conf)
-
-        # polling params from config file propagated to client object
-        with mock.patch("dwave.cloud.config.loaders.load_profile_from_files", lambda *pa, **kw: conf):
-            with dwave.cloud.Client.from_config() as client:
-                self.assertEqual(client.config.polling_schedule.backoff_min, 0.1)
-                self.assertEqual(client.config.polling_schedule.backoff_max, 1.0)
-
-        # test defaults
-        conf = dict(token='token')
-        with mock.patch("dwave.cloud.config.loaders.load_profile_from_files", lambda *pa, **kw: conf):
-            with dwave.cloud.Client.from_config() as client:
-                self.assertEqual(client.config.polling_schedule.backoff_min, Client.DEFAULTS['poll_backoff_min'])
-                self.assertEqual(client.config.polling_schedule.backoff_max, Client.DEFAULTS['poll_backoff_max'])
+        with self.subTest("defaults"):
+            conf = {"token": "token"}
+            with mock.patch("dwave.cloud.config.loaders.load_profile_from_files",
+                            lambda *pa, **kw: conf):
+                with dwave.cloud.Client.from_config() as client:
+                    ps = client.config.polling_schedule
+                    self.assertEqual(ps.strategy, Client.DEFAULTS['poll_strategy'])
+                    self.assertEqual(ps.backoff_min, Client.DEFAULTS['poll_backoff_min'])
+                    self.assertEqual(ps.backoff_max, Client.DEFAULTS['poll_backoff_max'])
+
+        with self.subTest("backoff"):
+            conf = {"poll_strategy": "backoff",
+                    "poll_backoff_min": "0.1",
+                    "poll_backoff_max": "1",
+                    "token": "token"}
+            with mock.patch("dwave.cloud.config.loaders.load_profile_from_files",
+                            lambda *pa, **kw: conf):
+                with dwave.cloud.Client.from_config() as client:
+                    ps = client.config.polling_schedule
+                    self.assertEqual(ps.strategy, PollingStrategy.BACKOFF)
+                    self.assertEqual(ps.backoff_min, 0.1)
+                    self.assertEqual(ps.backoff_max, 1.0)
+
+        with self.subTest("long polling"):
+            conf = {"poll_strategy": "long-polling",
+                    "poll_wait_time": "10",
+                    "poll_pause": "1",
+                    "token": "token"}
+            with mock.patch("dwave.cloud.config.loaders.load_profile_from_files", lambda *pa, **kw: conf):
+                with dwave.cloud.Client.from_config() as client:
+                    ps = client.config.polling_schedule
+                    self.assertEqual(ps.strategy, PollingStrategy.LONG_POLLING)
+                    self.assertEqual(ps.wait_time, 10)
+                    self.assertEqual(ps.pause, 1.0)
 
+    @mock.patch("dwave.cloud.config.loaders.load_profile_from_files",
+                lambda *pa, **kw: dict(token='token'))
     def test_polling_params_from_kwargs(self):
-        poll_conf = {"poll_backoff_min": "0.1", "poll_backoff_max": "1"}
-        conf = dict(token='token', **poll_conf)
-
-        with mock.patch("dwave.cloud.config.loaders.load_profile_from_files", lambda *pa, **kw: conf):
-            with dwave.cloud.Client.from_config(poll_backoff_min=0.5) as client:
+        with self.subTest("backoff"):
+            with dwave.cloud.Client.from_config(poll_strategy="backoff",
+                                                poll_backoff_min=0.5) as client:
                 self.assertEqual(client.config.polling_schedule.backoff_min, 0.5)
-                self.assertEqual(client.config.polling_schedule.backoff_max, 1.0)
+
+        with self.subTest("long polling"):
+            with dwave.cloud.Client.from_config(poll_strategy="long-polling",
+                                                poll_wait_time=10) as client:
+                self.assertEqual(client.config.polling_schedule.wait_time, 10)
 
     def _verify_retry_config(self, retry, opts):
         self.assertEqual(retry.total, opts['http_retry_total'])
diff --git a/tests/test_config.py b/tests/test_config.py
index 5f4f644..b072fea 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -29,7 +29,7 @@ from dwave.cloud.testing import iterable_mock_open
 from dwave.cloud.config import (
     get_configfile_paths, load_config_from_files, load_config,
     parse_float, parse_int, parse_boolean, get_cache_dir, update_config)
-from dwave.cloud.config.models import ClientConfig
+from dwave.cloud.config.models import ClientConfig, PollingStrategy
 from dwave.cloud.config.models import validate_config_v1, load_config_v1, dump_config_v1
 
 
@@ -644,6 +644,27 @@ class TestConfigModel(unittest.TestCase):
                      get_field=lambda config: getattr(config.polling_schedule, key),
                      model_value=model_value)
 
+    @parameterized.expand([
+        ("default",
+            {}, "strategy", PollingStrategy.BACKOFF),
+        ("explicit backoff",
+            {"poll_strategy": "backoff"}, "strategy", PollingStrategy.BACKOFF),
+        ("explicit long-polling",
+            {"poll_strategy": "long-polling"}, "strategy", PollingStrategy.LONG_POLLING),
+        ("implicit strategy, conformant value accepted",
+            {"poll_backoff_min": 1}, "backoff_min", 1.0),
+        ("implicit strategy, non-conformant value ignored",
+            {"poll_wait_time": 1}, "wait_time", OMITTED),
+        ("explicit strategy, conformant value accepted",
+            {"poll_wait_time": 1, "poll_strategy": "long-polling"}, "wait_time", 1.0),
+        ("explicit strategy, non-conformant value ignored",
+            {"poll_backoff_min": 1, "poll_strategy": "long-polling"}, "backoff_min", OMITTED),
+    ])
+    def test_polling_schedule_union(self, desc, raw_config, key, model_value):
+        self._verify(raw_config=raw_config,
+                     get_field=lambda config: getattr(config.polling_schedule, key, self.OMITTED),
+                     model_value=model_value)
+
     @parameterized.expand([
         ("1", True), ("on", True), ("On", True), ("true", True), ("True", True),
         ("0", False), ("off", False), ("Off", False), ("false", False), ("False", False),
diff --git a/tests/test_mock_submission.py b/tests/test_mock_submission.py
index bba0faf..355de85 100644
--- a/tests/test_mock_submission.py
+++ b/tests/test_mock_submission.py
@@ -20,7 +20,9 @@ import threading
 import time
 import unittest
 from concurrent.futures import TimeoutError, ThreadPoolExecutor
+from typing import Any, Dict
 from unittest import mock
+from urllib.parse import urlencode
 
 import numpy
 from parameterized import parameterized
@@ -82,7 +84,19 @@ def choose_reply(path, replies, statuses=None, date=None):
         raise NotImplementedError(path)
 
 
-class _QueryTest(unittest.TestCase):
+def fix_status_paths(replies: Dict[str, Any], **params):
+    """Append URL params to problem status paths."""
+
+    def fix(path, **params):
+        query = urlencode(params)
+        if 'problems/?id=' in path and query:
+            return f"{path}&{query}"
+        return path
+
+    return {fix(path, **params): reply for path, reply in replies.items()}
+
+
+class _QueryTest:
     def _check(self, results, linear, quad, offset=0, num_reads=1):
         # Did we get the right number of samples?
         self.assertEqual(num_reads, sum(results.num_occurrences))
@@ -112,12 +126,17 @@ class _QueryTest(unittest.TestCase):
 
 
 @mock.patch('time.sleep', lambda *x: None)
-class MockSubmission(_QueryTest):
+class MockSubmissionBaseTests(_QueryTest):
     """Test connecting and some related failure modes."""
 
     @classmethod
     def setUpClass(cls):
         cls.sapi = StructuredSapiMockResponses()
+        cls.config = dict(
+            endpoint='endpoint',
+            token='token',
+            poll_strategy=cls.poll_strategy,
+        )
 
     def test_submit_null_reply(self):
         """Get an error when the server's response is incomplete."""
@@ -131,7 +150,7 @@ class MockSubmission(_QueryTest):
             return session
 
         with mock.patch.object(Client, 'create_session', create_mock_session):
-            with Client(endpoint='endpoint', token='token') as client:
+            with Client(**self.config) as client:
                 solver = Solver(client, self.sapi.solver.data)
 
                 linear, quadratic = self.sapi.problem
@@ -154,7 +173,7 @@ class MockSubmission(_QueryTest):
             return session
 
         with mock.patch.object(Client, 'create_session', create_mock_session):
-            with Client(endpoint='endpoint', token='token') as client:
+            with Client(**self.config) as client:
                 solver = Solver(client, self.sapi.solver.data)
 
                 linear, quadratic = self.sapi.problem
@@ -178,7 +197,7 @@ class MockSubmission(_QueryTest):
             return session
 
         with mock.patch.object(Client, 'create_session', create_mock_session):
-            with Client(endpoint='endpoint', token='token') as client:
+            with Client(**self.config) as client:
                 solver = Solver(client, self.sapi.solver.data)
 
                 h, J = self.sapi.problem
@@ -211,7 +230,7 @@ class MockSubmission(_QueryTest):
             return session
 
         with mock.patch.object(Client, 'create_session', create_mock_session):
-            with Client(endpoint='endpoint', token='token') as client:
+            with Client(**self.config) as client:
                 solver = Solver(client, self.sapi.solver.data)
 
                 qubo = {(0, 0): 4.0, (0, 4): -4, (4, 4): 4.0}
@@ -247,7 +266,7 @@ class MockSubmission(_QueryTest):
             return session
 
         with mock.patch.object(Client, 'create_session', create_mock_session):
-            with Client(endpoint='endpoint', token='token') as client:
+            with Client(**self.config) as client:
                 solver = Solver(client, self.sapi.solver.data)
 
                 qubo = {(0, 0): 4.0, (0, 4): -4, (4, 4): 4.0}
@@ -272,7 +291,7 @@ class MockSubmission(_QueryTest):
             return session
 
         with mock.patch.object(Client, 'create_session', create_mock_session):
-            with Client(endpoint='endpoint', token='token') as client:
+            with Client(**self.config) as client:
                 solver = Solver(client, self.sapi.solver.data)
 
                 linear, quadratic = self.sapi.problem
@@ -294,7 +313,7 @@ class MockSubmission(_QueryTest):
             return session
 
         with mock.patch.object(Client, 'create_session', create_mock_session):
-            with Client(endpoint='endpoint', token='token') as client:
+            with Client(**self.config) as client:
                 solver = Solver(client, self.sapi.solver.data)
 
                 linear, quadratic = self.sapi.problem
@@ -320,7 +339,7 @@ class MockSubmission(_QueryTest):
             return session
 
         with mock.patch.object(Client, 'create_session', create_mock_session):
-            with Client(endpoint='endpoint', token='token') as client:
+            with Client(**self.config) as client:
                 solver = Solver(client, self.sapi.solver.data)
 
                 linear, quadratic = self.sapi.problem
@@ -350,7 +369,7 @@ class MockSubmission(_QueryTest):
             return session
 
         with mock.patch.object(Client, 'create_session', create_mock_session):
-            with Client(endpoint='endpoint', token='token') as client:
+            with Client(**self.config) as client:
                 solver = Solver(client, self.sapi.solver.data)
 
                 linear, quadratic = self.sapi.problem
@@ -371,14 +390,14 @@ class MockSubmission(_QueryTest):
             session.post = lambda a, _: choose_reply(a, {
                 'problems/': [self.sapi.continue_reply(id='123')]
             })
-            session.get = lambda a: choose_reply(a, {
+            session.get = lambda a: choose_reply(a, fix_status_paths({
                 'problems/?id=123': [self.sapi.complete_no_answer_reply(id='123')],
                 'problems/123/': self.sapi.complete_reply(id='123')
-            })
+            }, **self.poll_params))
             return session
 
         with mock.patch.object(Client, 'create_session', create_mock_session):
-            with Client(endpoint='endpoint', token='token') as client:
+            with Client(**self.config) as client:
                 solver = Solver(client, self.sapi.solver.data)
 
                 linear, quadratic = self.sapi.problem
@@ -396,12 +415,13 @@ class MockSubmission(_QueryTest):
             session = mock.Mock()
             session.post = lambda a, _: choose_reply(a, {
                 'problems/': [self.sapi.continue_reply(id='123')]})
-            session.get = lambda a: choose_reply(a, {
-                'problems/?id=123': [self.sapi.error_reply(id='123')]})
+            session.get = lambda a: choose_reply(a, fix_status_paths({
+                'problems/?id=123': [self.sapi.error_reply(id='123')]
+            }, **self.poll_params))
             return session
 
         with mock.patch.object(Client, 'create_session', create_mock_session):
-            with Client(endpoint='endpoint', token='token') as client:
+            with Client(**self.config) as client:
                 solver = Solver(client, self.sapi.solver.data)
 
                 linear, quadratic = self.sapi.problem
@@ -423,7 +443,7 @@ class MockSubmission(_QueryTest):
             def continue_then_complete(path, state={'count': 0}):
                 state['count'] += 1
                 if state['count'] < 2:
-                    return choose_reply(path, {
+                    return choose_reply(path, fix_status_paths({
                         'problems/?id=1': [self.sapi.continue_reply(id='1')],
                         'problems/?id=2': [self.sapi.continue_reply(id='2')],
                         'problems/1/': self.sapi.continue_reply(id='1'),
@@ -432,9 +452,9 @@ class MockSubmission(_QueryTest):
                                              self.sapi.continue_reply(id='2')],
                         'problems/?id=2,1': [self.sapi.continue_reply(id='2'),
                                              self.sapi.continue_reply(id='1')]
-                    })
+                    }, **self.poll_params))
                 else:
-                    return choose_reply(path, {
+                    return choose_reply(path, fix_status_paths({
                         'problems/?id=1': [self.sapi.error_reply(id='1')],
                         'problems/?id=2': [self.sapi.complete_no_answer_reply(id='2')],
                         'problems/1/': self.sapi.error_reply(id='1'),
@@ -443,7 +463,7 @@ class MockSubmission(_QueryTest):
                                              self.sapi.complete_no_answer_reply(id='2')],
                         'problems/?id=2,1': [self.sapi.complete_no_answer_reply(id='2'),
                                              self.sapi.error_reply(id='1')]
-                    })
+                    }, **self.poll_params))
 
             def accept_problems_with_continue_reply(path, body, ids=iter('12')):
                 problems = json.loads(body)
@@ -459,7 +479,7 @@ class MockSubmission(_QueryTest):
         session = global_mock_session()
 
         with mock.patch.object(Client, 'create_session', lambda self: session):
-            with Client(endpoint='endpoint', token='token') as client:
+            with Client(**self.config) as client:
                 solver = Solver(client, self.sapi.solver.data)
 
                 linear, quadratic = self.sapi.problem
@@ -472,6 +492,13 @@ class MockSubmission(_QueryTest):
                     self._check(results1, linear, quadratic, **params)
                 self._check(results2, linear, quadratic, **params)
 
+
+class MockSubmissionWithShortPolling(MockSubmissionBaseTests,
+                                     unittest.TestCase):
+
+    poll_strategy = "backoff"
+    poll_params = {}
+
     def test_exponential_backoff_polling(self):
         "After each poll, back-off should double"
 
@@ -489,15 +516,15 @@ class MockSubmission(_QueryTest):
             def continue_then_complete(path, state={'count': 0}):
                 state['count'] += 1
                 if state['count'] < 3:
-                    return choose_reply(path, {
+                    return choose_reply(path, fix_status_paths({
                         'problems/?id=123': [self.sapi.continue_reply(id='123')],
                         'problems/123/': self.sapi.continue_reply(id='123')
-                    })
+                    }, **self.poll_params))
                 else:
-                    return choose_reply(path, {
+                    return choose_reply(path, fix_status_paths({
                         'problems/?id=123': [self.sapi.complete_no_answer_reply(id='123')],
                         'problems/123/': self.sapi.complete_reply(id='123')
-                    })
+                    }, **self.poll_params))
 
             session.get = continue_then_complete
 
@@ -506,7 +533,7 @@ class MockSubmission(_QueryTest):
         session = global_mock_session()
 
         with mock.patch.object(Client, 'create_session', lambda self: session):
-            with Client(endpoint='endpoint', token='token') as client:
+            with Client(**self.config) as client:
                 solver = Solver(client, self.sapi.solver.data)
 
                 future = solver.sample_qubo({})
@@ -528,14 +555,14 @@ class MockSubmission(_QueryTest):
             session.post = lambda path, _: choose_reply(path, {
                 'problems/': [self.sapi.continue_reply(id='1')]
             })
-            session.get = lambda path: choose_reply(path, {
+            session.get = lambda path: choose_reply(path, fix_status_paths({
                 'problems/?id=1': [self.sapi.complete_no_answer_reply(id='1')],
                 'problems/1/': self.sapi.complete_reply(id='1')
-            })
+            }, **self.poll_params))
             return session
 
         with mock.patch.object(Client, 'create_session', create_mock_session):
-            with Client(endpoint='endpoint', token='token') as client:
+            with Client(**self.config) as client:
                 solver = Solver(client, self.sapi.solver.data)
 
                 def assert_no_delay(s):
@@ -559,14 +586,14 @@ class MockSubmission(_QueryTest):
             session.post = lambda path, _: choose_reply(path, {
                 'problems/': [self.sapi.continue_reply(id='1')]
             }, date=badnow)
-            session.get = lambda path: choose_reply(path, {
+            session.get = lambda path: choose_reply(path, fix_status_paths({
                 'problems/?id=1': [self.sapi.complete_no_answer_reply(id='1')],
                 'problems/1/': self.sapi.complete_reply(id='1')
-            }, date=badnow)
+            }, **self.poll_params), date=badnow)
             return session
 
         with mock.patch.object(Client, 'create_session', create_mock_session):
-            with Client(endpoint='endpoint', token='token') as client:
+            with Client(**self.config) as client:
                 solver = Solver(client, self.sapi.solver.data)
 
                 def assert_no_delay(s):
@@ -596,18 +623,18 @@ class MockSubmission(_QueryTest):
             def continue_then_complete(path, state={'count': 0}):
                 state['count'] += 1
                 if state['count'] < 3:
-                    return choose_reply(path, replies={
+                    return choose_reply(path, replies=fix_status_paths({
                         'problems/?id=123': [self.sapi.continue_reply(id='123')],
                         'problems/123/': self.sapi.continue_reply(id='123')
-                    }, statuses={
+                    }, **self.poll_params), statuses=fix_status_paths({
                         'problems/?id=123': statuses,
                         'problems/123/': statuses
-                    })
+                    }, **self.poll_params))
                 else:
-                    return choose_reply(path, {
+                    return choose_reply(path, fix_status_paths({
                         'problems/?id=123': [self.sapi.complete_no_answer_reply(id='123')],
                         'problems/123/': self.sapi.complete_reply(id='123')
-                    })
+                    }, **self.poll_params))
 
             session.get = continue_then_complete
 
@@ -616,7 +643,7 @@ class MockSubmission(_QueryTest):
         session = global_mock_session()
 
         with mock.patch.object(Client, 'create_session', lambda self: session):
-            with Client(endpoint='endpoint', token='token') as client:
+            with Client(**self.config) as client:
                 solver = Solver(client, self.sapi.solver.data)
 
                 future = solver.sample_qubo({})
@@ -629,6 +656,13 @@ class MockSubmission(_QueryTest):
                     schedule.backoff_min * schedule.backoff_base**2)
 
 
+class MockSubmissionWithLongPolling(MockSubmissionBaseTests,
+                                    unittest.TestCase):
+
+    poll_strategy = "long-polling"
+    poll_params = dict(timeout=Client.DEFAULTS['poll_wait_time'])
+
+
 class DeleteEvent(Exception):
     """Throws exception when mocked client submits an HTTP DELETE request."""