onflow / flow-go

A fast, secure, and developer-friendly blockchain built to support the next generation of games, apps, and the digital assets that power them.
GNU Affero General Public License v3.0
534 stars 179 forks source link

[Access] Implement web socket connection mock for unit testing of web socket controller #6635

Open Guitarheroua opened 2 weeks ago

illia-malachyn commented 2 weeks ago

The WebsocketController accepts a WebSocket connection as an argument. To effectively unit test this, we need to mock the WebSocket connection. However, the Gorilla WebSocket library's connection interface is extensive and challenging to mock directly. To simplify testing, we can create a wrapper interface around the Gorilla WebSocket connection. This wrapper will be easier to mock, allowing us to simulate the WebSocket's behavior with custom channels that facilitate reading and writing data.

The wrapper interface, WebsocketConnection, defines the core methods for testing: ReadJSON, WriteJSON, and Close. Using this interface, we can verify that our controller sends and receives the expected data and behaves as intended.

The code outline below shows how this can be implemented:

type WebsocketConnection interface {
    ReadJSON(v interface{}) error
    WriteJSON(v interface{}) error
    Close() error
}

type GorillaWebsocketConnection struct {
    conn *websocket.Conn
}

func NewGorillaWebsocketConnection(conn *websocket.Conn) *GorillaWebsocketConnection {
    return &GorillaWebsocketConnection{
        conn: conn,
    }
}

// Ensures that GorillaWebsocketConnection satisfies the WebsocketConnection interface
var _ WebsocketConnection = (*GorillaWebsocketConnection)(nil)

In our test setup, a handler can initialize a new WebSocket connection using NewGorillaWebsocketConnection and pass it to the controller:

newConn := NewGorillaWebsocketConnection(conn)
controller := NewWebSocketController(logger, h.websocketConfig, h.streamApi, h.streamConfig, newConn)
controller.HandleConnection(context.TODO())

To facilitate testing, a mock struct, WebsocketConnectionMock, is defined. This mock struct includes custom channels (e.g., socket and testSocket) for simulating the connection behavior, as well as a closed state to indicate if the connection has been terminated.

type WebsocketConnectionMock struct {
    mock.Mock
    socket     chan interface{}
    closed     bool
}

By using this setup, we can easily simulate data exchange through the mock channels and test that WebsocketController performs as expected.

illia-malachyn commented 2 weeks ago

This should be a P0