diff --git a/src/python/oftest/controller.py b/src/python/oftest/controller.py
index cc8b97c..0769437 100644
--- a/src/python/oftest/controller.py
+++ b/src/python/oftest/controller.py
@@ -104,6 +104,7 @@
         self.socs = []
         self.connect_cv = Condition()
         self.message_cv = Condition()
+        self.disconnect_cv = Condition()
 
         # Counters
         self.socket_errors = 0
@@ -218,8 +219,7 @@
                                   % (hdr.version, OFP_VERSION))
                 print "Version %d does not match OFTest version %d" % \
                     (hdr.version, OFP_VERSION)
-                self.active = False
-                self.switch_socket = None
+                self.disconnect()
                 return
 
             msg = of_message_parse(rawmsg)
@@ -359,25 +359,23 @@
         self.logger.info("Waiting for switch connection")
         self.socs = [self.listen_socket]
         self.dbg_state = "running"
+
         while self.active:
             try:
                 sel_in, sel_out, sel_err = \
                     select.select(self.socs, [], self.socs, 1)
             except:
                 print sys.exc_info()
-                self.logger.error("Select error, exiting")
-                self.active = False
-                break
+                self.logger.error("Select error, disconnecting")
+                self.disconnect()
 
             for s in sel_err:
-                self.logger.error("Got socket error on: " + str(s))
-                self.active = False
-                break
+                self.logger.error("Got socket error on: " + str(s) + ", disconnecting")
+                self.disconnect()
 
             for s in sel_in:
                 if self._socket_ready_handle(s) == -1:
-                    self.active = False
-                    break
+                    self.disconnect()
 
         # End of main loop
         self.dbg_state = "closing"
@@ -396,6 +394,30 @@
             timed_wait(self.connect_cv, lambda: self.switch_socket, timeout=timeout)
         return self.switch_socket is not None
         
+    def disconnect(self, timeout=-1):
+        """
+        If connected to a switch, disconnect.
+        """
+        if self.switch_socket:
+            self.socs.remove(self.switch_socket)
+            self.switch_socket.close()
+            self.switch_socket = None
+            self.switch_addr = None
+            with self.disconnect_cv:
+                self.disconnect_cv.notifyAll()
+
+    def wait_disconnected(self, timeout=-1):
+        """
+        @param timeout Block for up to timeout seconds. Pass -1 for the default.
+        @return Boolean, True if disconnected
+        """
+
+        with self.disconnect_cv:
+            timed_wait(self.disconnect_cv, 
+                       lambda: True if not self.switch_socket else None, 
+                       timeout=timeout)
+        return self.switch_socket is None
+        
     def kill(self):
         """
         Force the controller thread to quit