From 17118c84ba385afcc7079ea4163f650db64ea6a7 Mon Sep 17 00:00:00 2001
From: gadial <gadial@gmail.com>
Date: Thu, 22 Jun 2023 15:39:19 +0300
Subject: [PATCH] Revising the completion_times property

---
 qiskit_experiments/framework/experiment_data.py  | 16 ++++++++--------
 test/database_service/test_db_experiment_data.py | 13 +++++++++++++
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/qiskit_experiments/framework/experiment_data.py b/qiskit_experiments/framework/experiment_data.py
index a1c1ef23f6..8837e94a07 100644
--- a/qiskit_experiments/framework/experiment_data.py
+++ b/qiskit_experiments/framework/experiment_data.py
@@ -338,6 +338,7 @@ def __init__(
         self._result_data = ThreadSafeList()
         self._figures = ThreadSafeOrderedDict(self._db_data.figure_names)
         self._analysis_results = ThreadSafeOrderedDict()
+        self._completion_times = ThreadSafeOrderedDict()
 
         self._deleted_figures = deque()
         self._deleted_analysis_results = deque()
@@ -362,12 +363,7 @@ def experiment(self):
     @property
     def completion_times(self) -> Dict[str, datetime]:
         """Returns the completion times of the jobs."""
-        job_times = {}
-        for job_id, job in self._jobs.items():
-            if job is not None and "COMPLETED" in job.time_per_step():
-                job_times[job_id] = job.time_per_step().get("COMPLETED")
-
-        return job_times
+        return dict(self._completion_times)
 
     @property
     def tags(self) -> List[str]:
@@ -887,8 +883,11 @@ def _add_job_data(
             job_result = job.result()
             self._add_result_data(job_result, jid)
             LOG.debug("Job data added [Job ID: %s]", jid)
-            # sets the endtime to be the time the last successful job was added
-            self.end_datetime = datetime.now()
+            completed_at = datetime.now()
+            with self._completion_times.lock:
+                self._completion_times[jid] = completed_at
+                # sets the end time to be the time the last successful job was added
+                self.end_datetime = completed_at
             return jid, True
         except Exception as ex:  # pylint: disable=broad-except
             # Handle cancelled jobs
@@ -2275,6 +2274,7 @@ def __json_encode__(self):
             "_deleted_figures": self._deleted_figures,
             "_deleted_analysis_results": self._deleted_analysis_results,
             "_result_data": self._result_data,
+            "_completion_times": self._completion_times,
             "_extra_data": self._extra_data,
             "_created_in_db": self._created_in_db,
             "_figures": self._safe_serialize_figures(),  # Convert figures to SVG
diff --git a/test/database_service/test_db_experiment_data.py b/test/database_service/test_db_experiment_data.py
index ec7ee55a4b..14a58b98f5 100644
--- a/test/database_service/test_db_experiment_data.py
+++ b/test/database_service/test_db_experiment_data.py
@@ -724,6 +724,19 @@ def _analysis(*args):  # pylint: disable = unused-argument
         self.assertEqual(exp_data.analysis_status(), AnalysisStatus.CANCELLED)
         self.assertEqual(exp_data.status(), ExperimentStatus.CANCELLED)
 
+    def test_completion_times(self):
+        """Test the completion_times property"""
+        jid = "1234"
+        job = mock.create_autospec(Job, instance=True)
+        job.job_id.return_value = jid
+        job.status = JobStatus.DONE
+
+        exp_data = ExperimentData(experiment_type="qiskit_test")
+        exp_data.add_jobs(job)
+        completion_times = exp_data.completion_times
+        self.assertTrue(jid in completion_times)
+        self.assertTrue(isinstance(completion_times[jid], datetime))
+
     def test_partial_cancel_analysis(self):
         """Test canceling experiment analysis."""