diff --git a/tests/test_canister/src/test_canister/test_canister.did b/tests/test_canister/src/test_canister/test_canister.did
index d911f36..268eefe 100644
--- a/tests/test_canister/src/test_canister/test_canister.did
+++ b/tests/test_canister/src/test_canister/test_canister.did
@@ -1,84 +1,4 @@
-type ClientPublicKey = blob;
-
-type DirectClientMessage = record {
-  client_key : ClientPublicKey;
-  message : blob;
-};
-
-type RelayedClientMessage = record {
-  content : blob;
-  sig : blob;
-};
-
-type CanisterIncomingMessage = variant {
-  DirectlyFromClient : DirectClientMessage;
-  RelayedByGateway : RelayedClientMessage;
-  IcWebSocketEstablished : ClientPublicKey;
-};
-
-type CanisterOutputMessage = record {
-  client_key : ClientPublicKey;
-  key : text;
-  val : blob;
-};
-
-type CanisterOutputCertifiedMessages = record {
-  messages : vec CanisterOutputMessage;
-  cert : blob;
-  tree : blob;
-};
-
-type CanisterWsRegisterArguments = record {
-  client_key : ClientPublicKey;
-};
-
-type CanisterWsRegisterResult = variant {
-  Ok : null;
-  Err : text;
-};
-
-type CanisterWsOpenArguments = record {
-  msg : blob;
-  sig : blob;
-};
-
-type CanisterWsOpenResultValue = record {
-  client_key : ClientPublicKey;
-  canister_id : principal;
-  nonce : nat64;
-};
-
-type CanisterWsOpenResult = variant {
-  Ok : CanisterWsOpenResultValue;
-  Err : text;
-};
-
-type CanisterWsCloseArguments = record {
-  client_key : ClientPublicKey;
-};
-
-type CanisterWsCloseResult = variant {
-  Ok : null;
-  Err : text;
-};
-
-type CanisterWsMessageArguments = record {
-  msg : CanisterIncomingMessage;
-};
-
-type CanisterWsMessageResult = variant {
-  Ok : null;
-  Err : text;
-};
-
-type CanisterWsGetMessagesArguments = record {
-  nonce : nat64;
-};
-
-type CanisterWsGetMessagesResult = variant {
-  Ok : CanisterOutputCertifiedMessages;
-  Err : text;
-};
+import "./ws_types.did";
 
 type CanisterWsSendResult = variant {
   Ok : null;
diff --git a/tests/test_canister/src/test_canister/ws_types.did b/tests/test_canister/src/test_canister/ws_types.did
new file mode 100644
index 0000000..70c01f0
--- /dev/null
+++ b/tests/test_canister/src/test_canister/ws_types.did
@@ -0,0 +1,86 @@
+type ClientPublicKey = blob;
+
+type DirectClientMessage = record {
+  client_key : ClientPublicKey;
+  message : blob;
+};
+
+type RelayedClientMessage = record {
+  content : blob;
+  sig : blob;
+};
+
+type GatewayStatusMessage = record {
+  status_index : nat64;
+};
+
+type CanisterIncomingMessage = variant {
+  DirectlyFromClient : DirectClientMessage;
+  RelayedByGateway : RelayedClientMessage;
+  IcWebSocketEstablished : ClientPublicKey;
+  IcWebSocketGatewayStatus : GatewayStatusMessage;
+};
+
+type CanisterOutputMessage = record {
+  client_key : ClientPublicKey;
+  key : text;
+  val : blob;
+};
+
+type CanisterOutputCertifiedMessages = record {
+  messages : vec CanisterOutputMessage;
+  cert : blob;
+  tree : blob;
+};
+
+type CanisterWsRegisterArguments = record {
+  client_key : ClientPublicKey;
+};
+
+type CanisterWsRegisterResult = variant {
+  Ok : null;
+  Err : text;
+};
+
+type CanisterWsOpenArguments = record {
+  msg : blob;
+  sig : blob;
+};
+
+type CanisterWsOpenResultValue = record {
+  client_key : ClientPublicKey;
+  canister_id : principal;
+  nonce : nat64;
+};
+
+type CanisterWsOpenResult = variant {
+  Ok : CanisterWsOpenResultValue;
+  Err : text;
+};
+
+type CanisterWsCloseArguments = record {
+  client_key : ClientPublicKey;
+};
+
+type CanisterWsCloseResult = variant {
+  Ok : null;
+  Err : text;
+};
+
+type CanisterWsMessageArguments = record {
+  msg : CanisterIncomingMessage;
+};
+
+type CanisterWsMessageResult = variant {
+  Ok : null;
+  Err : text;
+};
+
+type CanisterWsGetMessagesArguments = record {
+  nonce : nat64;
+};
+
+type CanisterWsGetMessagesResult = variant {
+  Ok : CanisterOutputCertifiedMessages;
+  Err : text;
+};