Skip to content

Commit

Permalink
chore: test the close message logic
Browse files Browse the repository at this point in the history
  • Loading branch information
@ilbertt
ilbertt committed Dec 9, 2023
1 parent e49876d commit 28288a5
Show file tree
Hide file tree
Showing 4 changed files with 116 additions and 46 deletions.
101 changes: 79 additions & 22 deletions src/ic-websocket.test.ts
View file
Original file line number Diff line number Diff line change
@@ -1,20 +1,44 @@
import WsMockServer from "jest-websocket-mock";
import { rest } from "msw";
import { setupServer } from "msw/node";
import { CallRequest, Cbor, fromHex } from "@dfinity/agent";
import { CallRequest, Cbor } from "@dfinity/agent";
import { IDL } from "@dfinity/candid";
import { Principal } from "@dfinity/principal";

import IcWebSocket, { MAX_ALLOWED_NETWORK_LATENCY_MS, createWsConfig } from "./ic-websocket";
import { Principal } from "@dfinity/principal";
import { generateRandomIdentity } from "./identity";
import { CanisterWsMessageArguments, CanisterWsOpenArguments, ClientKey, WebsocketServiceMessageContent, _WS_CANISTER_SERVICE, decodeWebsocketServiceMessageContent, isClientKeyEq, wsMessageIdl, wsOpenIdl } from "./idl";
import {
CanisterWsMessageArguments,
CanisterWsOpenArguments,
ClientKey,
WebsocketServiceMessageContent,
_WS_CANISTER_SERVICE,
decodeWebsocketServiceMessageContent,
wsMessageIdl,
wsOpenIdl,
} from "./idl";
import { GatewayHandshakeMessage } from "./types";
import type { WsAgentRequestMessage } from "./agent/types";

import { canisterId, client1Key } from "./test/clients";
import { INVALID_HANDSHAKE_MESSAGE_FROM_GATEWAY, INVALID_MESSAGE_KEY, VALID_ACK_MESSAGE, VALID_HANDSHAKE_MESSAGE_FROM_GATEWAY, VALID_MESSAGE_SEQ_NUM_2, VALID_MESSAGE_SEQ_NUM_3, VALID_OPEN_MESSAGE, encodeHandshakeMessage } from "./test/messages";
import {
INVALID_HANDSHAKE_MESSAGE_FROM_GATEWAY,
INVALID_MESSAGE_KEY,
VALID_ACK_MESSAGE,
VALID_CLOSE_MESSAGE,
VALID_HANDSHAKE_MESSAGE_FROM_GATEWAY,
VALID_MESSAGE_SEQ_NUM_3,
VALID_OPEN_MESSAGE,
encodeHandshakeMessage,
} from "./test/messages";
import { sleep } from "./test/helpers";
import { getTestCanisterActor, getTestCanisterActorWithoutMethods, getTestCanisterActorWrongArgs, getTestCanisterActorWrongOpt } from "./test/actor";
import type { WsAgentRequestMessage } from "./agent/types";
import { GATEWAY_PRINCIPAL } from "./test/constants";
import { GatewayHandshakeMessage } from "./types";
import {
getTestCanisterActor,
getTestCanisterActorWithoutMethods,
getTestCanisterActorWrongArgs,
getTestCanisterActorWrongOpt,
} from "./test/actor";
import { GATEWAY_PRINCIPAL, LOCAL_REPLICA_ROOT_KEY } from "./test/constants";

const wsGatewayAddress = "ws://127.0.0.1:8080";
// the canister from which the application message was sent (needed to verify the message certificate)
Expand All @@ -38,7 +62,7 @@ const mockReplica = setupServer(
ctx.status(200),
// this response was generated from the same local replica
// used to generate the messages below
ctx.body(fromHex("d9d9f7a66e69635f6170695f76657273696f6e66302e31382e3068726f6f745f6b65795885308182301d060d2b0601040182dc7c0503010201060c2b0601040182dc7c05030201036100948a091fa3439c49aa8782da536348bba3a525cc0b63c0e202797ae7baf38f615e5375b694818b4a1a5b0fb07242aede15eb79f6454c19c1ee54fd8b9c14dbb06d94df2f2a3cc4f6336f0419680025f4411f0d764aa0b6e9fd246ba71a80fad66c696d706c5f76657273696f6e65302e382e3069696d706c5f68617368784030343366663064393237626337313431643761643630616235646331313934636364303164393761386431633333393632643236663730323461646463336135757265706c6963615f6865616c74685f737461747573676865616c746879706365727469666965645f686569676874181b")),
ctx.body(LOCAL_REPLICA_ROOT_KEY),
);
}),
);
Expand Down Expand Up @@ -314,16 +338,16 @@ describe("IcWebsocket class", () => {
// send the open confirmation message from the canister
mockWsServer.send(Cbor.encode(VALID_OPEN_MESSAGE));
await sleep(100);
expect(onMessage).not.toHaveBeenCalled();

// send the ack message that has sequence number 2
mockWsServer.send(Cbor.encode(VALID_ACK_MESSAGE));
await sleep(100);
expect(onMessage).not.toHaveBeenCalled();
expect(onError).not.toHaveBeenCalled();

// send an application message from the canister
mockWsServer.send(Cbor.encode(VALID_MESSAGE_SEQ_NUM_2));

// wait for the message to be processed
mockWsServer.send(Cbor.encode(VALID_MESSAGE_SEQ_NUM_3));
await sleep(100);

expect(onMessage).toHaveBeenCalled();
expect(onError).not.toHaveBeenCalled();
});
Expand All @@ -345,16 +369,12 @@ describe("IcWebsocket class", () => {
// send the open confirmation message from the canister
mockWsServer.send(Cbor.encode(VALID_OPEN_MESSAGE));
await sleep(100);

expect(onMessage).not.toHaveBeenCalled();
expect(onError).not.toHaveBeenCalled();

// send an application message from the canister
mockWsServer.send(Cbor.encode(VALID_MESSAGE_SEQ_NUM_3));

// wait for the message to be processed
await sleep(100);

expect(onMessage).not.toHaveBeenCalled();
const seqNumError = new Error("[onWsMessage] Received message sequence number does not match next expected value. Expected: 2, received: 3");
expect(onError).toHaveBeenCalledWith(new ErrorEvent("error", { error: new Error(`Error receiving message: ${seqNumError}`) }));
Expand All @@ -378,16 +398,12 @@ describe("IcWebsocket class", () => {
// send the open confirmation message from the canister
mockWsServer.send(Cbor.encode(VALID_OPEN_MESSAGE));
await sleep(100);

expect(onMessage).not.toHaveBeenCalled();
expect(onError).not.toHaveBeenCalled();

// send an application message from the canister
mockWsServer.send(Cbor.encode(INVALID_MESSAGE_KEY));

// wait for the message to be processed
await sleep(100);

expect(onMessage).not.toHaveBeenCalled();
const invalidCertificateError = new Error("[onWsMessage] Certificate validation failed");
expect(onError).toHaveBeenCalledWith(new ErrorEvent("error", { error: new Error(`Error receiving message: ${invalidCertificateError}`) }));
Expand Down Expand Up @@ -454,6 +470,47 @@ describe("IcWebsocket class", () => {
expect(IDL.decode([IDL.Record({ 'text': IDL.Text })], wsMessageArgs[0].msg.content as Uint8Array)[0]).toMatchObject(applicationMessageContent);
expect(wsMessageArgs[1]).toEqual([]); // check that we're not sending unneeded arguments
});

it("closes the connection if close message is received", async () => {
const onMessage = jest.fn();
const onError = jest.fn();
const onClose = jest.fn();
const icWs = new IcWebSocket(wsGatewayAddress, undefined, icWebsocketConfig);
expect(icWs).toBeDefined();
// workaround: simulate the client identity
icWs["_clientKey"] = client1Key;
icWs.onmessage = onMessage;
icWs.onerror = onError;
icWs.onclose = onClose;
await mockWsServer.connected;
await sendHandshakeMessage(VALID_HANDSHAKE_MESSAGE_FROM_GATEWAY);

// we need to send the previous messages in order
// to not break the sequence number count

// send the open confirmation message from the canister
mockWsServer.send(Cbor.encode(VALID_OPEN_MESSAGE));
await sleep(100);
expect(onClose).not.toHaveBeenCalled();

// send the ack message that has sequence number 2
mockWsServer.send(Cbor.encode(VALID_ACK_MESSAGE));
await sleep(100);
expect(onClose).not.toHaveBeenCalled();

// send an application message from the canister
mockWsServer.send(Cbor.encode(VALID_MESSAGE_SEQ_NUM_3));
await sleep(100);
expect(onClose).not.toHaveBeenCalled();

// finally, send the close message
mockWsServer.send(Cbor.encode(VALID_CLOSE_MESSAGE));
await sleep(100);
expect(onMessage).toHaveBeenCalledTimes(1); // only with the application message
expect(onError).not.toHaveBeenCalled();
expect(onClose).toHaveBeenCalled();
await expect(mockWsServer.closed).resolves.not.toThrow();
});
});

describe("Messages acknowledgement", () => {
Expand Down
6 changes: 2 additions & 4 deletions src/test/clients.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,7 @@ import { Principal } from "@dfinity/principal";

export const canisterId = Principal.fromText("bnz7o-iuaaa-aaaaa-qaaaa-cai");

// Principal: "pmisz-prtlk-b6oe6-bj4fl-6l5fy-h7c2h-so6i7-jiz2h-bgto7-piqfr-7ae"
// const client1Seed = "rabbit fun moral twin food kangaroo egg among adjust pottery measure seek";
export const client1Key: ClientKey = {
client_principal: Principal.fromText("pmisz-prtlk-b6oe6-bj4fl-6l5fy-h7c2h-so6i7-jiz2h-bgto7-piqfr-7ae"),
client_nonce: BigInt("5768810803147064100"),
client_principal: Principal.fromText("kj67s-b5v2y-ahlkr-kmume-xbow6-zwbtj-j4j3m-ae46e-qqrcu-uxiby-yae"),
client_nonce: BigInt("385892949151814926"),
};
5 changes: 4 additions & 1 deletion src/test/constants.ts
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
import { fromHex } from "@dfinity/agent";
import { Principal } from "@dfinity/principal";

export const GATEWAY_PRINCIPAL = Principal.fromText("i3gux-m3hwt-5mh2w-t7wwm-fwx5j-6z6ht-hxguo-t4rfw-qp24z-g5ivt-2qe");
export const GATEWAY_PRINCIPAL = Principal.fromText("sqdfl-mr4km-2hfjy-gajqo-xqvh7-hf4mf-nra4i-3it6l-neaw4-soolw-tae");

export const LOCAL_REPLICA_ROOT_KEY = fromHex("d9d9f7a66e69635f6170695f76657273696f6e66302e31382e3068726f6f745f6b65795885308182301d060d2b0601040182dc7c0503010201060c2b0601040182dc7c050302010361008005229d89a17c6f9ec403a4b1a8aa103fc48055046c95f1e60ee2fbfb0bb23ab21617a93f48b99b1199ac89008cf3cf0a83e9da35f5cf27d0d51535ceff89c43ee236c31c3a7865cc6b333194ad3f7155b2931a7ffec2066777dffb20f277ca6c696d706c5f76657273696f6e65302e382e3069696d706c5f68617368784064613931633732316637386462393433346561336630303437383939383836346439313731346538626561363862333963633736326662306263383937313662757265706c6963615f6865616c74685f737461747573676865616c746879706365727469666965645f68656967687418d4");
Loading

0 comments on commit 28288a5

Please sign in to comment.