Building a Real-Time Chat Gateway with NestJS: Testing WebSocket Features and Implementing Simple Private Messaging

Introduction

WebSocket technology has revolutionized real-time communication on the web, enabling bidirectional data flow between clients and servers. NestJS, a powerful Node.js framework, provides excellent support for building WebSocket applications using its WebSocket Gateways. In this guide, we will explore how to effectively test NestJS WebSocket Gateways, step by step.

In the first part of this series, we started by setting up a basic NestJS application and explored how to effectively test WebSocket gateways using Jest and Socket.IO.

In this second part, we will build on the foundation laid in the previous article and dive deeper into testing by implementing new features for our chat gateway. We will also explore more advanced testing techniques to ensure the reliability of our chat WebSocket gateway.

Let's get started!

Disclaimer: It's important to note that there is no need to follow the previous part of this series to understand and follow along with this article. However, if you are interested in the initial setup and testing of WebSocket gateways using Jest and Socket.IO, you can find the first part of the series here.

Part 1: Setting Up the Project and Dependencies

Next, we will clone a sample repository that we'll use as the starting point for our project. Use the following commands to clone the repository:

git clone https://github.com/jfrancai/nestjs-playground.git nest-app

And navigate into the project folder:

cd nest-app

Now, let's take a look at the folder structure of the project:

.
├── nest-cli.json
├── package.json
├── package-lock.json
├── README.md
├── src
│   ├── app.controller.spec.ts
│   ├── app.controller.ts
│   ├── app.module.ts
│   ├── app.service.ts
│   ├── chat
│   │   ├── chat.gateway.spec.ts
│   │   └── chat.gateway.ts
│   └── main.ts
├── test
│   ├── app.e2e-spec.ts
│   └── jest-e2e.json
├── tsconfig.build.json
└── tsconfig.json

In the src folder, you'll find various files related to the NestJS application, including the WebSocket Gateway for the chat module. The test folder contains end-to-end testing files (we are not going to use it for now), and the package.json file lists the project's dependencies.

Setting Up Your Own Git Repository

Since we cloned the sample repository, it is essential to set up your own Git repository for your project. We'll start by removing the current .git folder:

rm -rf .git

Next, initialize a new Git repository in the project folder:

git init
git add .
git commit -m 'first commit'

With this, you have your own Git repository ready to track your project's changes.

Install Node Dependencies

Now, let's install the necessary Node.js dependencies for our project. Run the following command to install them:

npm install

The package.json file contains the list of dependencies required for the project:

{
  "dependencies": {
    // ...
    "@nestjs/platform-socket.io": "^10.1.3",
    "@nestjs/websockets": "^10.1.3"
  },
  "devDependencies": {
    // ...
    "socket.io-client": "^4.7.2"
  }
}

Note that we have socket.io-client installed as a dev dependency, which will be used to test our application backend. Additionally, the jest configuration in the same file sets up the testing environment.

Jest Configuration

Jest is the default testing library provided with NestJS. The configuration for Jest can be found in the package.json file:

{
  "jest": {
    "moduleFileExtensions": ["js", "json", "ts"],
    "rootDir": "src",
    "testRegex": ".*\\.spec\\.ts$",
    "transform": {
      "^.+\\.(t|j)s$": "ts-jest"
    },
    "collectCoverageFrom": ["**/*.(t|j)s"],
    "coverageDirectory": "../coverage",
    "testEnvironment": "node"
  }
}

For more detailed information on Jest configuration options, you can refer to the official Jest documentation available here.

Starting Jest Watch Mode

Now that we have everything set up, we can start Jest in watch mode, which allows it to automatically rerun the test suite whenever a file change is detected. Open a new shell window and run the following command:

npm run test:watch

This command executes Jest in watch mode. If you look inside the package.json file, you'll find the actual command used jest --watch:

{
  // ...
  "scripts": {
    // ...
    "start:watch": "jest --watch"
  }
}

With Jest running in watch mode, you can now develop your NestJS WebSocket Gateway and have your tests automatically updated and executed on each change.

Part 2: Updating Existing Tests

Let's start by updating the test we wrote in the previous article to make it more readable and reusable. The test was responsible for checking if the server responds with "pong" when it receives a "ping" event. Here's the original test:

it('should emit "pong" on "ping"', async () => {
  ioClient.connect();
  ioClient.emit("ping", "Hello world!");
  await new Promise<void>((resolve) => {
    ioClient.on("connect", () => {
      console.log("connected");
    });
    ioClient.on("pong", (data) => {
      expect(data).toBe("Hello world!");
      resolve();
    });
  });
  ioClient.disconnect();
});

While this test works, it can be improved for better readability and reusability. To achieve this, we will create a helper function called eventReception, which will handle the asynchronous event reception during testing.

The eventReception Helper Function

Let's create the eventReception function inside the chat.gateway.spec.ts file:

async function eventReception(from: Socket, event: string): Promise<void> {
  return new Promise<void>((resolve) => {
    from.on(event, () => {
      resolve();
    });
  });
}

This function takes two parameters:

• from: A socket object (in this case, the client socket).
• event: The name of the event we expect the client to receive.

The function returns a new Promise, which will resolve when the specified event is received from the client. This allows us to await the reception of a specific event in our test.

Note: Before proceeding, let's briefly explain the concept of async/await and promises, which might have been covered too quickly in the previous article.

JavaScript Promise

A Promise is an object representing the eventual completion or failure of an asynchronous operation and its resulting value. It allows you to handle asynchronous operations in a more structured and readable way, avoiding the use of nested callbacks.

Understanding Async/Await

In JavaScript, asynchronous operations are common when dealing with tasks like fetching data from an API, reading files, or making network requests. Traditionally, asynchronous code was handled using callbacks, but it could lead to callback hell and make the code difficult to read and maintain.

To address these issues, ECMAScript 2017 (ES8) introduced the async and await keywords, which provide a more elegant way to write asynchronous code in a synchronous style. Here's a breakdown of how async and await work:

async Function

When you define a function with the async keyword, it becomes an asynchronous function. An asynchronous function always returns a Promise implicitly. The function can contain one or more asynchronous operations, such as API calls or file reads.

async function fetchData() {
  // Asynchronous operations can be performed here
  return result; // The function returns a Promise
}

await Operator

The await keyword is used inside an async function to wait for the resolution of a Promise. When you use await, the function execution will pause until the Promise is resolved or rejected. This allows you to handle asynchronous code in a more sequential manner, making it easier to read and write.

async function fetchData() {
  const result = await fetch("https://api.example.com/data");
  // Code here will wait until the fetch Promise is resolved
  return result;
}

Error Handling

Async functions can also use try...catch to handle errors from awaited promises. If the awaited Promise is rejected, the catch block will catch the error.

async function fetchData() {
  try {
    const result = await fetch("https://api.example.com/data");
    return result;
  } catch (error) {
    // Handle the error here
  }
}

The async and await keywords offer a more structured and readable way to work with asynchronous operations in JavaScript. By using async/await, you can avoid the complexities of nested callbacks and write asynchronous code in a style that resembles synchronous code. This feature is especially useful for writing clear and maintainable tests that involve asynchronous operations.

Rewriting the Test Using eventReception

Now that we have the eventReception helper function, let's rewrite our test:

it("should emit 'pong' on 'ping'", async () => {
  ioClient.on("connect", () => {
    console.log("connected");
  });

  ioClient.on("pong", (data) => {
    expect(data).toBe("Hello world!");
  });

  ioClient.connect();
  await eventReception(ioClient, "connect");

  ioClient.emit("ping", "Hello world!");
  await eventReception(ioClient, "pong");

  ioClient.disconnect();
});

With the eventReception function, our test becomes cleaner, more readable, and easily reusable in other tests.

In the next part, we will delve into implementing new features for our chat gateway. We'll enable user private messaging. Additionally, we'll emit an event to notify all clients when a new user joins the server.

Part 3: Implementing Private Messaging

Let's start by writing a test for the private messaging feature. The test will ensure that a message sent from one client socket (ioClient) is received by another client socket (receiverClient) as a private message. Here's the test:

it("should send a private message to another client", async () => {
  receiverClient.on("private message", (data) => {
    expect(data).toHaveProperty("from");
    expect(data.from).toBe(ioClient.id);
    expect(data).toHaveProperty("message");
    expect(data.message).toBe("Hello from the other side");
  });

  ioClient.connect();
  await eventReception(ioClient, "connect");

  receiverClient.connect();
  await eventReception(receiverClient, "connect");

  ioClient.emit("private message", {
    from: ioClient.id,
    to: receiverClient.id,
    message: "Hello from the other side",
  });
  await eventReception(receiverClient, "private message");

  ioClient.disconnect();
  receiverClient.disconnect();
});

1. We create an event listener on the receiverClient for the "private message" event. The test checks if the received data object contains properties for "from" and "message", and if the values are as expected.

2. We connect both ioClient and receiverClient to the chat gateway using ioClient.connect() and receiverClient.connect() respectively. We then use the eventReception helper function to wait for the "connect" event to be emitted on both clients before proceeding with the test.

3. Next, we use ioClient.emit("private message", ...) to send a private message from ioClient to receiverClient. The message object contains "from", "to", and "message" properties.

4. Finally, we use eventReception(receiverClient, "private message") to wait for the "private message" event to be received by receiverClient.

The test ensures that the private message is correctly sent and received between the two clients.

Note: In the test setup, the receiverClient is declared in the beforeEach section of the test suite like for ioClient. The beforeEach function is a hook provided by Jest, and it is executed before each test case in the current test suite.

describe('ChatGateway', () => {
  //...
  let receiverClient: Socket;

  beforeEach(async () => {
    //...
    receiverClient = io('http://localhost:3000', {
      autoConnect: false,
      transports: ['websocket', 'polling'],
    });
});

Implementing Private Messaging in the Chat Gateway

To pass the test, we need to add the code that handles private messages in our ChatGateway class controller.

Add the following code to the ChatGateway class controller in src/chat/chat.gateway.ts:

@SubscribeMessage('private message')
handlePrivateMessage(client: any, data: any) {
    client.to(data.to).emit('private message', {
      from: client.id,
      message: data.message,
    });
}

The handlePrivateMessage function is a WebSocket event handler for the "private message" event. When the server receives a "private message" event from a client, it will call this function.

Inside the function, we use the client.to(data.to).emit(...) method to send the private message to the specified recipient (data.to). The message object contains the sender's ID (from) and the message text (message).

The test we wrote earlier ensures that private messages are correctly sent and received between clients.

Part 4: Start Using Promise.all()

In this part, we will introduce a new feature for our chat gateway. We want clients to receive information about the list of already connected clients upon their own connection to the server. We will achieve this by emitting a custom event called 'connected clients' to each newly connected client.

Writing the Test

Let's start by writing a test to ensure that the server sends the correct list of connected clients to a newly connected client.

We create an event listener on the receiverClient for the "connected clients" event. When the server emits this event, the test will check the data received.

We utilize Jest assertions to ensure that the received data is an array with a length of two. Each element in the array should be an object containing a "userID" property with the corresponding ID of the ioClient and receiverClient, respectively.

it("should send the number of clients on the server when connecting", async () => {
  receiverClient.on("connected clients", (data) => {
    expect(data.length).toBe(2);
    expect(data).toEqual(
      expect.arrayContaining([
        expect.objectContaining({
          userID: ioClient.id,
        }),
        expect.objectContaining({
          userID: receiverClient.id,
        }),
      ]),
    );
  });

  // ...
});

We connect both ioClient and receiverClient to the chat gateway using ioClient.connect() and receiverClient.connect() respectively. We then use the eventReception helper function to wait for the "connect" event to be emitted on both clients before proceeding with the test.

it("should send the number of clients on the server when connecting", async () => {
  //...

  ioClient.connect();
  await eventReception(ioClient, "connect");

  receiverClient.connect();
  await eventReception(receiverClient, "connect");

  //...
});

Next, we use eventReception(receiverClient, "connected clients") to wait for the "connected clients" event to be received by receiverClient. This ensures that the server has sent the list of connected clients to the newly connected client.

it("should send the number of clients on the server when connecting", async () => {
  //...

  await eventReception(receiverClient, "connected clients");

  // And we don't forget to disconnect clients once the test is finishing
  ioClient.disconnect();
  receiverClient.disconnect();
});

The test ensures that the server correctly sends the list of connected clients to each client upon connection.

Implementing the "connected clients" Feature

Now that the test has been written, we need to implement the "connected clients" feature in our ChatGateway class controller.

Update the handleConnection function in src/chat/chat.gateway.ts:

handleConnection(client: any, ...args: any[]) {
  const { sockets } = this.io.sockets;
  this.logger.log(`Client id: ${client.id} connected`);
  this.logger.debug(`Number of connected clients: ${sockets.size}`);

  const connectedClients = Array.from(sockets).map(([_, socket]) => ({
    userID: socket.id,
  }));

  client.emit("connected clients", connectedClients);
}

In the handleConnection function, we first collect all the sockets connected to the server using this.io.sockets. Then, we create an array of objects containing the userID of each connected client.

Finally, we emit the "connected clients" event to the newly connected client (client) with the list of connected clients (connectedClients).

And the test should pass...

...But this approach may not always work as expected, depending on the situation.

In some cases, using the following code:

await eventReception(receiverClient, "connect");

followed by:

await eventReception(receiverClient, "connected clients");

may cause your test to pass or fail based on the delay between the "connect" event and the "connected clients" event. If the "connected clients" event arrives too quickly after the "connect" event, it may result in a test timeout on the second await eventReception call.

To address this issue, there are two approaches we can consider:

Approach 1: Removing the await eventReception(receiverClient, "connect") line, as we have already tested the "connect" event in a previous test. However, this is more of a workaround and does not solve the underlying problem. Additionally, in other scenarios, we might want to test the arrival of multiple messages, not knowing which one arrives first.

Approach 2: Using Promise.all

Promise.all is a built-in JavaScript function that takes an array of promises and returns a new promise that resolves when all the promises in the array have resolved, or rejects if any of the promises reject.

Here's how we can use Promise.all to handle multiple awaited events:

await Promise.all([
  eventReception(receiverClient, "connect"),
  eventReception(receiverClient, "connected clients"),
]);

With Promise.all, both eventReception promises will be executed simultaneously, and the test will wait for both events to be received before proceeding to the next step.

By using Promise.all, we can maintain a clean and expressive writing style for our tests while ensuring that we handle multiple awaited events efficiently and reliably.

Conclusion

In this article, we delved into enhancing our NestJS WebSocket gateway by implementing new features and writing effective tests using Jest and Socket.IO. We started by introducing private messaging, allowing clients to send messages directly to each other. We then explored how to emit an event containing the list of already connected clients to newly connected clients.

Throughout the process, we leveraged Jest's powerful assertion functions to validate the behavior of our chat gateway and ensure that it functions as intended. Additionally, we utilized async/await to handle asynchronous operations in a more synchronous and readable manner, making our tests more manageable and organized.

We also encountered challenges, such as handling multiple awaited events, which we skillfully resolved using Promise.all. This allowed us to efficiently synchronize and manage multiple asynchronous events in our tests, providing a robust testing framework for future enhancements.

In conclusion, I hope that this step-by-step guide has simplified the process of building a NestJS WebSocket gateway and provided valuable insights into effective testing strategies.

If you enjoyed reading this article and found it helpful, please consider giving it a thumbs-up or leaving a comment below. Your feedback and engagement encourage me to create more content like this, helping the community learn and grow together.

Happy coding 🚀🌟