diff --git a/testclient_test.go b/testclient_test.go
index beac45a..c54eb6d 100644
--- a/testclient_test.go
+++ b/testclient_test.go
@@ -315,6 +315,7 @@ func (c *TestClient) WaitForClientRemoved(ctx context.Context) error {
 		c.hub.mu.Unlock()
 		select {
 		case <-ctx.Done():
+			c.hub.mu.Lock()
 			return ctx.Err()
 		default:
 			time.Sleep(time.Millisecond)
@@ -341,6 +342,7 @@ func (c *TestClient) WaitForSessionRemoved(ctx context.Context, sessionId string
 		c.hub.mu.Unlock()
 		select {
 		case <-ctx.Done():
+			c.hub.mu.Lock()
 			return ctx.Err()
 		default:
 			time.Sleep(time.Millisecond)