Do not build a Job out of a result that is not a job id

tests/test_wait_non_job.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""JSONRPCClient.wait() with `job` set, for calls that never became a trackable job.
+
+With new-style jobs the server does not answer a job method with its id: the id arrives in a
+`core.get_jobs` event, which `_process_message` uses to set `Call.job_id`. A method that is not a
+job, and a job declared `transient=True`, produce no such event, so the call returns the method's
+own result and there is no job to wait on.
+"""
+import unittest
+from collections import defaultdict
+from threading import Event, Lock
+
+from truenas_api_client import Call, JSONRPCClient, Job
+from truenas_api_client.exc import ClientException
+
+
+def _client(new_style=True):
+    """Build a client without running __init__, which would open a connection."""
+    obj = object.__new__(JSONRPCClient)
+    obj._calls = {}
+    obj._jobs = defaultdict(dict)
+    obj._jobs_lock = Lock()
+    obj._call_timeout = 10
+    obj._new_style_jobs = new_style
+    return obj
+
+
+def _returned(method='some.method', result=None, job_id=None):
+    call = Call(method, ())
+    call.result = result
+    call.job_id = job_id
+    call.returned.set()
+    return call
+
+
+class TestWaitWithoutTrackableJob(unittest.TestCase):
+    """No `core.get_jobs` event named the call, so its own result is the answer."""
+
+    def test_non_job_method_returns_its_result(self):
+        client = _client()
+        call = _returned('vm.start', result=True)
+
+        self.assertIs(client.wait(call, job=True), True)
+
+    def test_transient_job_returns_its_result(self):
+        client = _client()
+        call = _returned('pool.scrub', result=None)
+
+        self.assertIsNone(client.wait(call, job=True))
+
+    def test_int_result_is_not_mistaken_for_a_job_id(self):
+        """The result of a non-job method can be an integer without being a job id."""
+        client = _client()
+        call = _returned('some.count', result=7)
+
+        self.assertEqual(client.wait(call, job=True), 7)
+
+    def test_job_return_raises_rather_than_waiting_forever(self):
+        client = _client()
+        call = _returned('vm.start', result=True)
+
+        with self.assertRaises(ClientException):
+            client.wait(call, job='RETURN')
+
+    def test_call_is_unregistered(self):
+        client = _client()
+        call = _returned('vm.start', result=True)
+        client._calls[call.id] = call
+
+        client.wait(call, job=True)
+
+        self.assertNotIn(call.id, client._calls)
+
+
+class TestWaitWithTrackableJob(unittest.TestCase):
+    """A `core.get_jobs` event set `job_id`, so the job is waited on as before."""
+
+    def test_uses_job_id_not_result(self):
+        client = _client()
+        call = _returned('pool.import_pool', result=42, job_id=42)
+        client._jobs[42].update(state='SUCCESS', result='done', __ready=Event())
+        client._jobs[42]['__ready'].set()
+
+        self.assertEqual(client.wait(call, job=True), 'done')
+
+    def test_job_return_gives_the_job(self):
+        client = _client()
+        call = _returned('pool.import_pool', result=42, job_id=42)
+
+        jobobj = client.wait(call, job='RETURN')
+
+        self.assertIsInstance(jobobj, Job)
+        self.assertEqual(jobobj.job_id, 42)
+
+
+class TestLegacyJobs(unittest.TestCase):
+    """Against a server without new-style jobs the id is the call's plain result."""
+
+    def test_result_is_used_as_the_job_id(self):
+        client = _client(new_style=False)
+        call = _returned('pool.import_pool', result=42)
+
+        jobobj = client.wait(call, job='RETURN')
+
+        self.assertIsInstance(jobobj, Job)
+        self.assertEqual(jobobj.job_id, 42)
+
+
+if __name__ == '__main__':
+    unittest.main()

truenas_api_client/__init__.py

@@ -325,6 +325,10 @@ def __init__(self, method: str, params: tuple):
         self.method = method
         self.params = params
         self.returned = Event()
+        self.job_id: Any = None
+        """Set when a `core.get_jobs` event binds this call to a job. Stays
+        `None` for a method that is not a job, and for a transient job, which
+        emits no such events."""
         self.result: Any = None
         self.error: ClientException | None = None
         self.py_exception: BaseException | None = None
@@ -558,6 +562,7 @@ def _recv(self, message: JSONRPCMessage):
                             if params['collection'] == 'core.get_jobs' and params['msg'] in ['added', 'changed']:
                                 for message_id in params['fields']['message_ids']:
                                     if (call := self._calls.get(message_id)) is not None:
+                                        call.job_id = params['id']
                                         call.result = params['id']
                                         call.returned.set()
                                         self._unregister_call(call)
@@ -891,11 +896,14 @@ def wait(
 
         Returns:
             Job: If `job='RETURN'`, return the `Job` object.
-            Any: If `job=True`, return the job's result. Otherwise, return the call's result.
+            Any: If `job=True`, return the job's result. Otherwise, return the call's result. A method that did
+                not run as a trackable job, because it is not a job or because it is transient, has already
+                returned its own result, which is returned as-is.
 
         Raises:
             CallTimeout: The call took longer than `timeout` seconds to return.
-            ClientException: The call ended in error and `py_exception` was not enabled for `c`.
+            ClientException: The call ended in error and `py_exception` was not enabled for `c`, or `job='RETURN'`
+                was requested for a method that did not run as a trackable job.
             BaseException: The call ended in error and `py_exception` was enabled for `c`.
 
         """
@@ -916,7 +924,29 @@ def wait(
                     raise c.error
 
             if job:
-                jobobj = Job(self, c.result, callback=callback)
+                if self._new_style_jobs:
+                    if c.job_id is None:
+                        # No `core.get_jobs` event ever named this call, so
+                        # there is no job to track: either the method is not a
+                        # job, or it is a transient job, which emits no such
+                        # events. Either way the call has already returned the
+                        # method's own result, and building a Job here would
+                        # wait forever on events for an id the server never
+                        # issued.
+                        if job == 'RETURN':
+                            raise ClientException(
+                                f'{c.method!r} did not run as a trackable job, so there is no job to return.'
+                            )
+
+                        return c.result
+
+                    job_id = c.job_id
+                else:
+                    # Legacy jobs: the server answers a job method with the job
+                    # id as the call's plain result.
+                    job_id = c.result
+
+                jobobj = Job(self, job_id, callback=callback)
                 if job == 'RETURN':
                     return jobobj
                 return jobobj.result()