Fix connection cleanup when context is canceled

Signed-off-by: Luca Bigliardi <shammash@google.com>

client/connection.go

@@ -503,13 +503,17 @@ func (conn *Conn) ping(ctx context.Context) {
 // It pulls Lines from the input channel and dispatches them to any
 // handlers that have been registered for that IRC verb.
 func (conn *Conn) runLoop(ctx context.Context) {
-	defer conn.wg.Done()
 	for {
 		select {
 		case line := <-conn.in:
 			conn.dispatch(line)
 		case <-ctx.Done():
-			// control channel closed, bail out
+			// control channel closed, trigger Cancel() to clean
+			// things up properly and bail out
+
+			// We can't defer this, because Close() waits for it.
+			conn.wg.Done()
+			conn.Close()
 			return
 		}
 	}

client/connection_test.go

@@ -103,6 +103,28 @@ func TestEOF(t *testing.T) {
 	}
 }
 
+func TestCleanupOnContextDone(t *testing.T) {
+	c, s := setUp(t)
+	// Since we're not using tearDown() here, manually call Finish()
+	defer s.ctrl.Finish()
+
+	// Close() triggers DISCONNECT handler after cleaning up the state
+	// use this as a proxy to check that Close() was indeed called
+	dcon := callCheck(t)
+	c.Handle(DISCONNECTED, dcon)
+
+	// Simulate context cancelation using our cancel func
+	c.die()
+
+	// Verify that disconnected handler was called
+	dcon.assertWasCalled("Conn did not call disconnected handlers.")
+
+	// Verify that the connection no longer thinks it's connected
+	if c.Connected() {
+		t.Errorf("Conn still thinks it's connected to the server.")
+	}
+}
+
 func TestClientAndStateTracking(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	st := state.NewMockTracker(ctrl)