Fixed the unit tests for the full websocket implementation

We landed the websocket implementation in steps, and asynchronously. This patch reflects the work needed to integrate all the pieces.
Tiny fix

A previous patch missed this change.