diff --git a/src/biokbase/narrative/jobs/job.py b/src/biokbase/narrative/jobs/job.py index dfd76ac891..2dd1a308d7 100644 --- a/src/biokbase/narrative/jobs/job.py +++ b/src/biokbase/narrative/jobs/job.py @@ -104,7 +104,7 @@ def __init__(self, ee2_state, extra_data=None, children=None): if ee2_state.get("job_id") is None: raise ValueError("Cannot create a job without a job ID!") - self._update_state(ee2_state) + self.update_state(ee2_state) self.extra_data = extra_data # verify parent-children relationship @@ -130,9 +130,9 @@ def from_job_ids(cls, job_ids, return_list=True): return jobs @staticmethod - def _trim_ee2_state(state: dict, exclude: list) -> None: - if exclude: - for field in exclude: + def _trim_ee2_state(state: dict, exclude_fields: list) -> None: + if exclude_fields: + for field in exclude_fields: if field in state: del state[field] @@ -241,7 +241,7 @@ def __getattr__(self, name): def __setattr__(self, name, value): if name in STATE_ATTRS: # TODO are/should these assignments be used? - self._update_state({name: value}) + self.update_state({name: value}) else: object.__setattr__(self, name, value) @@ -308,7 +308,7 @@ def parameters(self): f"Unable to fetch parameters for job {self.job_id} - {e}" ) - def _update_state(self, state: dict, ts: int = None) -> None: + def update_state(self, state: dict, ts: int = None) -> None: """ Given a state data structure (as emitted by ee2), update the stored state in the job object All updates to the job state should go through here to keep the last_updated field accurate @@ -320,7 +320,7 @@ def _update_state(self, state: dict, ts: int = None) -> None: if self._acc_state: if "job_id" in state and state["job_id"] != self.job_id: raise ValueError( - f"Job ID mismatch in _update_state: job ID: {self.job_id}; state ID: {state['job_id']}" + f"Job ID mismatch in update_state: job ID: {self.job_id}; state ID: {state['job_id']}" ) # Check if there would be no change in updating @@ -337,28 +337,28 @@ def _update_state(self, state: dict, ts: int = None) -> None: self.last_updated = time_ns() if ts is None else ts - def refresh_state(self, force_refresh=False, exclude=JOB_INIT_EXCLUDED_JOB_STATE_FIELDS): + def refresh_state(self, force_refresh=False, exclude_fields=JOB_INIT_EXCLUDED_JOB_STATE_FIELDS): """ Queries the job service to see the state of the current job. """ if force_refresh or not self.was_terminal(): state = self.query_ee2_state(self.job_id, init=False) - self._update_state(state) + self.update_state(state) - return self.current_state(exclude) + return self.cached_state(exclude_fields) - def current_state(self, exclude=None): + def cached_state(self, exclude_fields=None): """Wrapper for self._acc_state""" state = copy.deepcopy(self._acc_state) - self._trim_ee2_state(state, exclude) + self._trim_ee2_state(state, exclude_fields) return state def output_state(self, state=None, no_refresh=False) -> dict: """ :param state: Supplied when the state is queried beforehand from EE2 in bulk, or when it is retrieved from a cache. If not supplied, must be - queried with self.refresh_state() or self.current_state() + queried with self.refresh_state() or self.cached_state() :return: dict, with structure { @@ -411,10 +411,10 @@ def output_state(self, state=None, no_refresh=False) -> dict: :rtype: dict """ if not state: - state = self.current_state() if no_refresh else self.refresh_state() + state = self.cached_state() if no_refresh else self.refresh_state() else: - self._update_state(state) - state = self.current_state() + self.update_state(state) + state = self.cached_state() if state is None: return self._create_error_state( @@ -464,8 +464,8 @@ def show_output_widget(self, state=None): if not state: state = self.refresh_state() else: - self._update_state(state) - state = self.current_state() + self.update_state(state) + state = self.cached_state() if state["status"] == COMPLETED_STATUS and "job_output" in state: (output_widget, widget_params) = self._get_output_info(state) diff --git a/src/biokbase/narrative/jobs/jobmanager.py b/src/biokbase/narrative/jobs/jobmanager.py index 1dae531d03..6a1bffba93 100644 --- a/src/biokbase/narrative/jobs/jobmanager.py +++ b/src/biokbase/narrative/jobs/jobmanager.py @@ -318,7 +318,7 @@ def _construct_job_output_state_set( # pre-emptively try a job state update # so can mark the set of fetched (but also changed) states # with a simultaneous timestamp - job._update_state(fetched_state, now) + job.update_state(fetched_state, now) output_states[job_id] = job.output_state(fetched_state) else: # fetch the current state without updating it diff --git a/src/biokbase/narrative/tests/test_job.py b/src/biokbase/narrative/tests/test_job.py index 908e8977b9..966148f7a1 100644 --- a/src/biokbase/narrative/tests/test_job.py +++ b/src/biokbase/narrative/tests/test_job.py @@ -389,10 +389,10 @@ def test_job_update__no_state(self): # should fail with error 'state must be a dict' with self.assertRaisesRegex(TypeError, "state must be a dict"): - job._update_state(None) + job.update_state(None) self.assertFalse(job.was_terminal()) - job._update_state({}) + job.update_state({}) self.assertFalse(job.was_terminal()) @mock.patch(CLIENTS, get_mock_client) @@ -405,8 +405,8 @@ def test_job_update__invalid_job_id(self): self.assertEqual(job.refresh_state(), expected) # try to update it with the job state from a different job - with self.assertRaisesRegex(ValueError, "Job ID mismatch in _update_state"): - job._update_state(get_test_job(JOB_COMPLETED)) + with self.assertRaisesRegex(ValueError, "Job ID mismatch in update_state"): + job.update_state(get_test_job(JOB_COMPLETED)) @mock.patch(CLIENTS, get_mock_client) def test_job_update__last_updated__no_change(self): @@ -418,22 +418,22 @@ def test_job_update__last_updated__no_change(self): ee2_state = get_test_job(job_id) job._acc_state = get_test_job(job_id) - job._update_state(ee2_state) + job.update_state(ee2_state) self.assertEqual(last_updated, job.last_updated) trim_ee2_state(ee2_state, JOB_INIT_EXCLUDED_JOB_STATE_FIELDS) - job._update_state(ee2_state) + job.update_state(ee2_state) self.assertEqual(last_updated, job.last_updated) trim_ee2_state(ee2_state, EXCLUDED_JOB_STATE_FIELDS) - job._update_state(ee2_state) + job.update_state(ee2_state) self.assertEqual(last_updated, job.last_updated) trim_ee2_state(ee2_state, OUTPUT_STATE_EXCLUDED_JOB_STATE_FIELDS) - job._update_state(ee2_state) + job.update_state(ee2_state) self.assertEqual(last_updated, job.last_updated) - job._update_state({}) + job.update_state({}) self.assertEqual(last_updated, job.last_updated) # job has init ee2 state @@ -442,18 +442,18 @@ def test_job_update__last_updated__no_change(self): trim_ee2_state(ee2_state, JOB_INIT_EXCLUDED_JOB_STATE_FIELDS) trim_ee2_state(job._acc_state, JOB_INIT_EXCLUDED_JOB_STATE_FIELDS) - job._update_state(ee2_state) + job.update_state(ee2_state) self.assertEqual(last_updated, job.last_updated) trim_ee2_state(ee2_state, EXCLUDED_JOB_STATE_FIELDS) - job._update_state(ee2_state) + job.update_state(ee2_state) self.assertEqual(last_updated, job.last_updated) trim_ee2_state(ee2_state, OUTPUT_STATE_EXCLUDED_JOB_STATE_FIELDS) - job._update_state(ee2_state) + job.update_state(ee2_state) self.assertEqual(last_updated, job.last_updated) - job._update_state({}) + job.update_state({}) self.assertEqual(last_updated, job.last_updated) @mock.patch(CLIENTS, get_mock_client) @@ -467,7 +467,7 @@ def test_job_update__last_updated__change(self): trim_ee2_state(job._acc_state, JOB_INIT_EXCLUDED_JOB_STATE_FIELDS) ee2_state = get_test_job(job_id) - job._update_state(ee2_state) + job.update_state(ee2_state) self.assertTrue(last_updated < job.last_updated) @mock.patch(CLIENTS, get_mock_client)