     (throw (Exception. (str "Unknown state: " state))))
   (let [message-state-atom (:message-state-atom broker-wrapper)
         message-id (message-id-key msg)]
-    (swap! message-state-atom #(assoc % message-id state))))
+    (swap! message-state-atom #(assoc % message-id state))
+    state))
 
 (defn handler-wrapper
   "Wraps handler function to count acks, retries, fails"

      (queue/reset queue-broker)
      (reset! message-state-atom {})
      (reset! id-sequence-atom 0)
      (finally
        (reset! resetting?-atom false)))))

(defn wait-for-indexing
  "Wait for all messages to be marked as processed"
  [broker-wrapper]
  (wait-for-states broker-wrapper [:processed]))

;; TODO
#_(def valid-message-modes
  "A list of the modes in which the message queue broker can operate.
  :normal - message functions are called and processed as they normally would.
  :retry - all messages return a retry when they are processed."
  #{:normal :retry})

;; TODO
#_(defn set-message-mode
  "Used to toggle message queue between normal mode and failure mode. In normal mode messages are
  processed normally. In failure mode all messages return failures. The failure mode is useful for
  automated tests verifying specific behaviors on failure."
  [mode]
  (if (valid-message-modes mode)
    ;; Use an atom to set state?
    "TODO"
    (throw (Exception. (str "Invalid message queue mode: " mode)))))


(defn create-queue-broker-wrapper

      ;; if in :normal mode do the same thing as before
      ;; if in :failure mode return
      ;; {:status :retry :message "Queue is in failure mode"}
      (let [resp (handler context msg)
            message-state (case (:status resp)
                            :ok (set-message-state broker-wrapper msg :processed)

                            :retry (when (queue/retry-limit-met? msg (count (iconfig/rabbit-mq-ttls)))
                                     (set-message-state broker-wrapper msg :failed))

                            :fail (set-message-state broker-wrapper msg :failed)

                            (throw (Exception. (str "Invalid response: " (pr-str resp)))))]
        (update-message-queue-history :process msg message-state)
        resp))))

(defn- current-message-states
