From f4ca04a9280727aee1fd575d67cab41c21bfa1e0 Mon Sep 17 00:00:00 2001
From: Dylan Anthony <43723790+dbanty@users.noreply.github.com>
Date: Wed, 25 Aug 2021 13:38:35 -0600
Subject: [PATCH] Fix: Compatibility with Qiskit's `run_circuits` and
 `CircuitSampler` utilities. (#19)

---
 qiskit_rigetti/_qcs_job.py | 9 ++++++++-
 tests/test_qcs_backend.py  | 1 -
 tests/test_qcs_job.py      | 6 ++++--
 3 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/qiskit_rigetti/_qcs_job.py b/qiskit_rigetti/_qcs_job.py
index 5843d9f..969b0f4 100644
--- a/qiskit_rigetti/_qcs_job.py
+++ b/qiskit_rigetti/_qcs_job.py
@@ -177,7 +177,14 @@ def cancel(self) -> None:
         raise NotImplementedError("Cancelling jobs is not supported")
 
     def status(self) -> JobStatus:
-        """Get the current status of this Job"""
+        """Get the current status of this Job
+
+        If this job was RUNNING when you called it, this function will block until the job is complete.
+        """
+
+        if self._status == JobStatus.RUNNING:
+            # Wait for results _now_ to finish running, otherwise consuming code might wait forever.
+            self.result()
 
         return self._status
 
diff --git a/tests/test_qcs_backend.py b/tests/test_qcs_backend.py
index 56e58e9..14a4162 100644
--- a/tests/test_qcs_backend.py
+++ b/tests/test_qcs_backend.py
@@ -27,7 +27,6 @@ def test_run(backend: RigettiQCSBackend):
     job = execute(circuit, backend, shots=10)
 
     assert job.backend() is backend
-    assert job.status() == JobStatus.RUNNING
     result = job.result()
     assert job.status() == JobStatus.DONE
     assert result.backend_name == backend.configuration().backend_name
diff --git a/tests/test_qcs_job.py b/tests/test_qcs_job.py
index b982204..9f2177a 100644
--- a/tests/test_qcs_job.py
+++ b/tests/test_qcs_job.py
@@ -171,12 +171,14 @@ def test_init__circuit_with_barrier(backend: RigettiQCSBackend, mocker: MockerFi
 
 
 def test_result(job: RigettiQCSJob):
+    assert job._status == JobStatus.RUNNING
+    assert job.status() == JobStatus.DONE, "Checking status did not wait for completion"
+    assert job._status == JobStatus.DONE
+
     result = job.result()
 
     assert result.date == job.result().date, "Result not cached"
 
-    assert job.status() == JobStatus.DONE
-
     assert result.backend_name == "3q-qvm"
     assert result.job_id == job.job_id()
     assert result.success is True