diff --git a/src/cubeb_audiounit.cpp b/src/cubeb_audiounit.cpp
index 25355b439..ad443cc59 100644
--- a/src/cubeb_audiounit.cpp
+++ b/src/cubeb_audiounit.cpp
@@ -498,8 +498,8 @@ audiounit_input_callback(void * user_ptr,
   assert(stm->input_unit != NULL);
   assert(AU_IN_BUS == bus);
 
-  if (stm->shutdown) {
-    ALOG("(%p) input shutdown", stm);
+  if (stm->shutdown || stm->reinit_pending) {
+    ALOG("(%p) input shutdown or reinit exit callback early.", stm);
     return noErr;
   }
 
@@ -605,20 +605,8 @@ audiounit_output_callback(void * user_ptr,
   long input_frames = 0;
   void * output_buffer = NULL, * input_buffer = NULL;
 
-  if (stm->shutdown) {
-    ALOG("(%p) output shutdown.", stm);
-    audiounit_make_silent(&outBufferList->mBuffers[0]);
-    return noErr;
-  }
-
-  if (stm->draining) {
-    OSStatus r = AudioOutputUnitStop(stm->output_unit);
-    assert(r == 0);
-    if (stm->input_unit) {
-      r = AudioOutputUnitStop(stm->input_unit);
-      assert(r == 0);
-    }
-    stm->state_callback(stm, stm->user_ptr, CUBEB_STATE_DRAINED);
+  if (stm->shutdown || stm->reinit_pending) {
+    ALOG("(%p) output shutdown or reinit exit callback early.", stm);
     audiounit_make_silent(&outBufferList->mBuffers[0]);
     return noErr;
   }
@@ -723,6 +711,12 @@ audiounit_output_callback(void * user_ptr,
     }
   }
 
+  if (stm->draining) {
+    audiounit_stream_stop_internal(stm);
+    stm->state_callback(stm, stm->user_ptr, CUBEB_STATE_DRAINED);
+    stm->shutdown = true;
+  }
+
   return noErr;
 }
 
diff --git a/tools/cubeb-test.cpp b/tools/cubeb-test.cpp
index e16a9bb01..dde034e5d 100644
--- a/tools/cubeb-test.cpp
+++ b/tools/cubeb-test.cpp
@@ -49,8 +49,8 @@ class cubeb_client final {
 
   bool init(char const * backend_name = nullptr);
   bool init_stream();
-  bool start_stream() const;
-  bool stop_stream() const;
+  bool start_stream();
+  bool stop_stream();
   bool destroy_stream() const;
   bool destroy();
   bool activate_log(cubeb_log_level log_level) const;
@@ -70,6 +70,8 @@ class cubeb_client final {
   cubeb_stream_params output_params = {};
   cubeb_stream_params input_params = {};
 
+  void force_drain() { _force_drain = true; }
+
 private:
   bool has_input() { return input_params.rate != 0; }
   bool has_output() { return output_params.rate != 0; }
@@ -85,6 +87,7 @@ class cubeb_client final {
   std::atomic<uint32_t> _channels = {0};
   std::atomic<bool> _latency_testing = {false};
   std::atomic<uint32_t> _latency_frames = {0}; // if !0, override. Else, use min.
+  std::atomic<bool> _force_drain = {false};
 
 
   /* Accessed only from audio thread. */
@@ -159,7 +162,8 @@ bool cubeb_client::init_stream() {
   return true;
 }
 
-bool cubeb_client::start_stream() const {
+bool cubeb_client::start_stream() {
+  _force_drain = false;
   int rv = cubeb_stream_start(stream);
   if (rv != CUBEB_OK) {
     fprintf(stderr, "Could not start the stream\n");
@@ -168,7 +172,8 @@ bool cubeb_client::start_stream() const {
   return true;
 }
 
-bool cubeb_client::stop_stream() const {
+bool cubeb_client::stop_stream() {
+  _force_drain = false;
   int rv = cubeb_stream_stop(stream);
   if (rv != CUBEB_OK) {
     fprintf(stderr, "Could not stop the stream\n");
@@ -275,6 +280,11 @@ long cubeb_client::user_data_cb(cubeb_stream* stm, void* user,
   }
 
   _total_frames += nframes;
+
+  if (_force_drain) {
+    return nframes - 1;
+  }
+
   return nframes;
 }
 
@@ -333,6 +343,7 @@ void print_help() {
     "2: change log level to verbose\n"
     "p: start a initialized stream\n"
     "s: stop a started stream\n"
+    "d: force stream to drain\n"
     "c: get stream position (client thread)\n"
     "i: change device type to input\n"
     "o: change device type to output\n"
@@ -345,7 +356,7 @@ void print_help() {
   fprintf(stderr, "%s\n", msg);
 }
 
-bool choose_action(const cubeb_client& cl, operation_data * op, int c) {
+bool choose_action(cubeb_client& cl, operation_data * op, int c) {
   // Consume "enter" and "space"
   while (c == 10 || c == 32) {
     c = getchar();
@@ -398,6 +409,8 @@ bool choose_action(const cubeb_client& cl, operation_data * op, int c) {
     } else {
       fprintf(stderr, "stop_stream failed\n");
     }
+  } else if (c == 'd') {
+    cl.force_drain();
   } else if (c == 'c') {
     uint64_t pos = cl.get_stream_position();
     uint64_t latency = cl.get_stream_latency();