Fix result shots (#70)

qiskit_braket_provider/providers/braket_job.py

@@ -24,7 +24,7 @@ def retry_if_result_none(result):
     wrap_exception=True,
 )
 def _get_result_from_aws_tasks(
-    tasks: Union[List[LocalQuantumTask], List[AwsQuantumTask]], shots: int
+    tasks: Union[List[LocalQuantumTask], List[AwsQuantumTask]]
 ) -> Optional[List[ExperimentResult]]:
     """Returns experiment results of AWS tasks.
 
@@ -55,7 +55,7 @@ def _get_result_from_aws_tasks(
             return None
 
         experiment_result = ExperimentResult(
-            shots=shots,
+            shots=result.task_metadata.shots,
             success=True,
             status=task.state(),
             data=data,
@@ -107,9 +107,7 @@ def submit(self):
         return
 
     def result(self) -> Result:
-        experiment_results = _get_result_from_aws_tasks(
-            tasks=self._tasks, shots=self.shots
-        )
+        experiment_results = _get_result_from_aws_tasks(tasks=self._tasks)
         return Result(
             backend_name=self._backend,
             backend_version=self._backend.version,

tests/providers/mocks.py

@@ -114,14 +114,14 @@
 
 MOCK_GATE_MODEL_QUANTUM_TASK_RESULT = GateModelQuantumTaskResult(
     task_metadata=TaskMetadata(
-        **{"id": str(uuid.uuid4()), "deviceId": "default", "shots": 2}
+        **{"id": str(uuid.uuid4()), "deviceId": "default", "shots": 3}
     ),
     additional_metadata=None,
-    measurements=[np.array([0, 1]), np.array([1, 0])],
+    measurements=np.array([[0, 1], [0, 1], [1, 0]]),
     measured_qubits=[0, 1],
     result_types=None,
     values=None,
-    measurement_counts=Counter({"01": 1, "10": 1}),
+    measurement_counts=Counter({"01": 2, "10": 1}),
 )
 
 MOCK_LOCAL_QUANTUM_TASK = LocalQuantumTask(MOCK_GATE_MODEL_QUANTUM_TASK_RESULT)

tests/providers/test_braket_job.py

@@ -1,6 +1,6 @@
 """Tests for AWS Braket job."""
 
-from unittest import TestCase, expectedFailure
+from unittest import TestCase
 
 from qiskit.providers import JobStatus
 
@@ -33,18 +33,8 @@ def test_result(self):
         job = self._get_job()
 
         self.assertEqual(job.result().job_id, "AwesomeId")
-        self.assertEqual(job.result().results[0].data.counts, {"01": 1, "10": 1})
-        self.assertEqual(job.result().results[0].data.memory, ["10", "01"])
+        self.assertEqual(job.result().results[0].data.counts, {"01": 1, "10": 2})
+        self.assertEqual(job.result().results[0].data.memory, ["10", "10", "01"])
         self.assertEqual(job.result().results[0].status, "COMPLETED")
-        self.assertEqual(job.result().get_memory(), ["10", "01"])
-
-    @expectedFailure
-    def test_result_shots(self):
-        """
-        Test result shots.
-
-        Expected to fail because shots is not extracted from the job.tasks.
-        """
-        job = self._get_job()
-
-        self.assertEqual(job.result().results[0].shots, 2)
+        self.assertEqual(job.result().results[0].shots, 3)
+        self.assertEqual(job.result().get_memory(), ["10", "10", "01"])