Repository: dwavesystems/dwave-cloud-client
Pull request #657 (2024-08-12 09:21:56)

Add `timeout` to `Problems.get_problem()`; enable long polling
I wasn't aware `timeout` is supported on get problem endpoint at the time I worked on #649.


Patch:
diff --git a/dwave/cloud/api/resources.py b/dwave/cloud/api/resources.py
index c9f4266..2686f60 100644
--- a/dwave/cloud/api/resources.py
+++ b/dwave/cloud/api/resources.py
@@ -179,10 +179,16 @@ class Problems(ResourceBase):
         return TypeAdapter(List[models.ProblemStatus]).validate_python(statuses)
 
     @accepts(media_type='application/vnd.dwave.sapi.problem+json', version='>=2.1,<3')
-    def get_problem(self, problem_id: str) -> models.ProblemStatusMaybeWithAnswer:
+    def get_problem(self,
+                    problem_id: str,
+                    timeout: Optional[int] = None,
+                    ) -> models.ProblemStatusMaybeWithAnswer:
         """Retrieve problem short status and answer if answer is available."""
         path = '{}'.format(problem_id)
-        response = self.session.get(path)
+        params = {}
+        if timeout is not None:
+            params.setdefault('timeout', timeout)
+        response = self.session.get(path, params=params)
         status = orjson.loads(response.content)
         return models.ProblemStatusMaybeWithAnswer.model_validate(status)
 
diff --git a/releasenotes/notes/get_problem-add-long-polling-timeout-8fb2f295cace01ad.yaml b/releasenotes/notes/get_problem-add-long-polling-timeout-8fb2f295cace01ad.yaml
new file mode 100644
index 0000000..0177a8e
--- /dev/null
+++ b/releasenotes/notes/get_problem-add-long-polling-timeout-8fb2f295cace01ad.yaml
@@ -0,0 +1,11 @@
+---
+features:
+  - |
+    Add ``timeout`` keyword argument to ``api.resources.Problems.get_problem()``
+    method.
+
+    Use it to enable long polling on problem fetch API endpoint. The connection
+    is closed (server-side) when the problem completes, or the ``timeout``
+    (in seconds) elapses, whatever comes first. In case problem completed in
+    time, problem answer is returned in addition to problem status (which is
+    always returned).


Test patch:
diff --git a/tests/api/resources/test_problems.py b/tests/api/resources/test_problems.py
index 038299c..b485d01 100644
--- a/tests/api/resources/test_problems.py
+++ b/tests/api/resources/test_problems.py
@@ -151,6 +151,18 @@ class ProblemResourcesBaseTests(abc.ABC):
         self.assertIsInstance(status.answer, models.ProblemAnswer)
         self.verify_problem_answer(status.answer)
 
+    def test_get_problem_long_polling(self):
+        """Problem status with answer retrieved by problem id, long polling variant"""
+
+        problem_id = self.problem_id
+        status = self.api.get_problem(problem_id, timeout=self.poll_wait_time)
+
+        self.assertIsInstance(status, models.ProblemStatusMaybeWithAnswer)
+        self.verify_problem_status(status, solved=True)
+
+        self.assertIsInstance(status.answer, models.ProblemAnswer)
+        self.verify_problem_answer(status.answer)
+
     def test_get_problem_status(self):
         """Problem status retrieved by problem id."""
 
@@ -454,6 +466,12 @@ class ProblemResourcesMockerMixin:
         # problem status et al
         self.mocker.get(
             url(f'problems/{p1_id}'),
+            complete_qs=True,
+            json=p1_status_with_answer,
+            request_headers=headers)
+        self.mocker.get(
+            url(f'problems/{p1_id}?timeout={self.poll_wait_time}'),
+            complete_qs=True,
             json=p1_status_with_answer,
             request_headers=headers)
         self.mocker.get(
@@ -478,18 +496,22 @@ class ProblemResourcesMockerMixin:
             request_headers=headers)
         self.mocker.get(
             url(f'problems/{p1_id}/info'),
+            complete_qs=True,
             json=p1_info,
             request_headers=headers)
         self.mocker.get(
             url(f'problems/{p1_id}/answer'),
+            complete_qs=True,
             json=p1_answer,
             request_headers=headers)
         self.mocker.get(
             p1_answer_data_uri,
+            complete_qs=True,
             body=p1_answer_data,
             request_headers=headers)
         self.mocker.get(
             url(f'problems/{p1_id}/messages'),
+            complete_qs=True,
             json=p1_messages,
             request_headers=headers)
         self.mocker.get(
@@ -516,10 +538,12 @@ class ProblemResourcesMockerMixin:
             request_headers=headers)
         self.mocker.get(
             url('problems/?max_results=1'),
+            complete_qs=True,
             json=[p1_status],
             request_headers=headers)
         self.mocker.get(
             url('problems/?status=COMPLETED'),
+            complete_qs=True,
             json=[p1_status],
             request_headers=headers)