Change async API to be more conductive to testing

src/metabase/feature_extraction/async.clj

@@ -35,12 +35,13 @@
         :status   :done
         :ended_at (u/new-sql-timestamp))))
   (swap! running-jobs dissoc id)
-  (log/info (format "Async job %s done." id)))
+  (log/info (format "Async job %s done." id))
+  payload)
 
 (defn- save-error
   [{:keys [id]} error]
-  (when-not (future-cancelled? (@running-jobs id))
-    (let [error (Throwable->map error)]
+  (let [error (Throwable->map error)]
+    (when-not (future-cancelled? (@running-jobs id))
       (log/warn (format "Async job %s encountered an error:\n%s." id error))
       (db/transaction
         (db/insert! ComputationJobResult
@@ -49,8 +50,9 @@
           :payload    error)
         (db/update! ComputationJob id
           :status :error
-          :ended_at (u/new-sql-timestamp)))))
-  (swap! running-jobs dissoc id))
+          :ended_at (u/new-sql-timestamp))))
+    (swap! running-jobs dissoc id)
+    error))
 
 (defn cancel
   "Cancel computation job (if still running)."

test/metabase/test/async.clj

@@ -2,7 +2,8 @@
   "Utilities for testing async API endpoints."
   (:require [clojure.tools.logging :as log]
             [metabase.feature-extraction.async :as async]
-            [metabase.models.computation-job :refer [ComputationJob]]))
+            [metabase.models.computation-job :refer [ComputationJob]]
+            [metabase.util :as u]))
 
 (def ^:dynamic ^Integer *max-while-runtime*
   "Maximal time in milliseconds `while-with-timeout` runs."
@@ -22,13 +23,16 @@
 (defn result!
   "Blocking version of async/result."
   [job-id]
-  (when-let [f (-> #'async/running-jobs
-                   deref                  ; var
-                   deref                  ; atom
-                   (get job-id))]
-    (when-not (or (future-cancelled? f)
-                  (future-done? f))
-      @f))
-  ; Wait for the transaction to finish
-  (while-with-timeout (-> job-id ComputationJob async/running?))
-  (async/result (ComputationJob job-id)))
+  (let [f (-> #'async/running-jobs
+              deref                  ; var
+              deref                  ; atom
+              (get job-id))]
+    (if (and f (not (future-cancelled? f)))
+      {:result     @f
+       :status     (-> job-id ComputationJob :status)
+       :created-at (u/new-sql-timestamp)}
+      (do
+        ;; Make sure the transaction has finished
+        (binding [*max-while-runtime* 1000]
+          (while-with-timeout (-> job-id ComputationJob async/running?)))
+        (async/result (ComputationJob job-id))))))