 (defn get-written-ledgers [onyx-client job-data onyx-id]
   (onyx.plugin.bookkeeper/read-ledgers-data (:log onyx-client) 
                                             onyx-id 
-                                            (:job-id @job-data)
-                                            (get-in @job-data [:task-ids :persist :id])))
-
-(defrecord WriteLogClient [client job-data ledger-handles ledger-handle ledger-ids onyx-client]
+                                            (:job-id job-data)
+                                            (get-in job-data [:task-ids :persist :id])))
+
+
+(defn close-ledger-handles [ledger-handles]
+  (mapv (fn [h] 
+          (try (.close h)
+               (catch Throwable t
+                 (info "Error closing handle" t)))) 
+        ledger-handles))
+
+(defn await-jobs-completions [peer-config jobs-data]
+  (mapv (fn [job-data]
+          (onyx.api/await-job-completion peer-config (:job-id job-data)))
+        jobs-data)) 
+
+(defrecord WriteLogClient [client jobs-data ledger-handles ledger-handle ledger-ids onyx-client]
   client/Client
   (setup! [this test node]
     (let [client (bookkeeper-client)

                                                 (try
                                                   (assoc op 
                                                          :type :ok 
                                                          :value (do
                                                                   (mapv (fn [h] 
                                                                           (try (.close h)
                                                                                (catch Throwable t
                                                                                  (info "Error closing handle" t)))) 
                                                                         @ledger-handles)
                                                                   (onyx.api/await-job-completion peer-config (:job-id @job-data))))
                                                   (catch Throwable t
                                                     (assoc op :type :info :value t))))
        :read-ledgers (timeout 1000000

(defn basic-test
  "A simple test of Onyx's safety."
  [version]
  ;; TODO, have map with test parameters in edn
  (let [n-jobs 1] 
    (merge tests/noop-test
           {:os os
            :db (setup version)
            :client (write-log-client (atom []) (atom []) (atom []))
            :model (->OnyxModel) ;; Currently not actually used for anything
            :checker (onyx-checker/->Checker peer-config (:n-peers peer-setup))
            :generator (gen/phases
                         (->> (onyx-gen/filter-new identity 
                                                   (onyx-gen/frequency [(adds) 
                                                                        (submit-job-gen n-jobs)
                                                                        ;(gen/once (gc-peer-logs))
                                                                        ]
                                                                       [0.98
                                                                        ;0.01
                                                                        0.01]))
                              (gen/stagger 1/10)
                              ;(gen/delay 1)
                              (gen/nemesis
                                (gen/seq (cycle
                                           [(gen/sleep 30)
                                            {:type :info :f :start}
                                            (gen/sleep 200)
                                            {:type :info :f :stop}])))
                              (gen/time-limit 14000)) 

                         ;; Bring everything back at the end
                         ;; This way we can test that the peers came back up
                         (gen/nemesis (gen/once {:type :info :f :stop}))
                         ;; Sleep for a while to give peers a chance to come back up
                         ;; Should be enough time that curator backoff * max-retries is covered
                         (gen/sleep 120)

                         (close-await-completion-gen)
                         (read-ledgers-gen)
                         (read-peer-log-gen))
            ;:nemesis (nemesis/partitioner (comp nemesis/bridge shuffle))
            :nemesis (nemesis/partition-random-halves)
            })))
