Support websocket connection callback function

__tests__/event_manager.test.ts

@@ -2,6 +2,7 @@
 import Ain from '../src/ain';
 import { FAILED_TO_REGISTER_ERROR_CODE } from '../src/constants';
 import { FilterDeletionReasons, TransactionStates } from '../src/types';
+import { sleep } from './test_util';
 const {
   test_node_3,
   test_event_handler_node,
@@ -10,11 +11,27 @@ const {
 jest.setTimeout(180000);
 
 describe('Event Handler', function() {
-  let ain = new Ain(test_node_3, test_event_handler_node);
+  const ain = new Ain(test_node_3, test_event_handler_node);
   let eventFilterId: string;
+  let connectionCount = 0;
+  let disconnectionCount = 0;
+  // Connection callback
+  const connectionCb = (websocket) => {
+    connectionCount++;
+  };
+  // Disconnection callback
+  const disconnectionCb = (websocket) => {
+    disconnectionCount++;
+  };
 
   beforeAll(async () => {
-    await ain.em.connect();
+    expect(connectionCount).toBe(0);
+    expect(disconnectionCount).toBe(0);
+
+    await ain.em.connect(connectionCb, disconnectionCb);
+
+    expect(connectionCount).toBe(1);
+    expect(disconnectionCount).toBe(0);
   });
 
   afterEach(async () => {
@@ -25,8 +42,15 @@ describe('Event Handler', function() {
     });
   })
 
-  afterAll(() => {
+  afterAll(async () => {
+    expect(connectionCount).toBe(1);
+    expect(disconnectionCount).toBe(0);
+
     ain.em.disconnect();
+
+    await sleep(3000);
+    expect(connectionCount).toBe(1);
+    expect(disconnectionCount).toBe(1);
   });
 
   describe('Channel connection', () => {

__tests__/test_util.ts

@@ -14,3 +14,8 @@ export function eraseStateVersion(result) {
   return erased;
 }
 
+export function sleep(ms) {
+  return new Promise((resolve) => {
+    setTimeout(resolve, ms);
+  });
+};

src/event-manager/event-channel-client.ts

@@ -6,6 +6,7 @@ import {
   EventChannelMessageTypes,
   EventChannelMessage,
   BlockchainEventTypes,
+  ConnectionCallback,
   DisconnectionCallback,
 } from '../types';
 import EventFilter from './event-filter';
@@ -51,10 +52,11 @@ export default class EventChannelClient {
 
   /**
    * Opens a new event channel.
+   * @param {ConnectionCallback} connectionCallback The connection callback function.
    * @param {DisconnectionCallback} disconnectionCallback The disconnection callback function.
    * @returns {Promise<void>} A promise for the connection success.
    */
-  connect(disconnectionCallback?: DisconnectionCallback): Promise<any> {
+  connect(connectionCallback?: ConnectionCallback, disconnectionCallback?: DisconnectionCallback): Promise<any> {
     return new Promise(async (resolve, reject) => {
       if (this.isConnected) {
         reject(new Error(`Can't connect multiple channels`));
@@ -137,11 +139,16 @@ export default class EventChannelClient {
         }
         // Heartbeat timeout
         this.startHeartbeatTimer(DEFAULT_HEARTBEAT_INTERVAL_MS);
+        // Connection callback
+        if (connectionCallback) {
+          connectionCallback(this._ws);
+        }
         resolve(this);
       };
 
       this._ws.onclose = () => {
         this.disconnect();
+        // Disconnection callback
         if (disconnectionCallback) {
           disconnectionCallback(this._ws);
         }

src/event-manager/index.ts

@@ -1,10 +1,18 @@
 import Ain from '../ain';
 import {
-  BlockFinalizedEventConfig, BlockFinalizedEvent,
+  BlockFinalizedEventConfig,
+  BlockFinalizedEvent,
   ErrorFirstCallback,
-  BlockchainEventConfig, BlockchainEventCallback,
-  TxStateChangedEventConfig, TxStateChangedEvent,
-  ValueChangedEventConfig, ValueChangedEvent, DisconnectionCallback, FilterDeletedEventCallback, BlockchainErrorCallback,
+  BlockchainEventConfig,
+  BlockchainEventCallback,
+  TxStateChangedEventConfig,
+  TxStateChangedEvent,
+  ValueChangedEventConfig,
+  ValueChangedEvent,
+  ConnectionCallback,
+  DisconnectionCallback,
+  FilterDeletedEventCallback,
+  BlockchainErrorCallback,
 } from '../types';
 import EventChannelClient from './event-channel-client';
 import EventCallbackManager from './event-callback-manager';
@@ -37,10 +45,11 @@ export default class EventManager {
 
   /**
    * Opens a new event channel.
+   * @param {ConnectionCallback} ConnectionCallback The connection callback function.
    * @param {DisconnectionCallback} disconnectionCallback The disconnection callback function.
    */
-  async connect(disconnectionCallback?: DisconnectionCallback) {
-    await this._eventChannelClient.connect(disconnectionCallback);
+  async connect(connectionCallback?: ConnectionCallback, disconnectionCallback?: DisconnectionCallback) {
+    await this._eventChannelClient.connect(connectionCallback, disconnectionCallback);
   }
 
   /**

src/types.ts

@@ -494,6 +494,11 @@ export type BlockchainErrorCallback = (error: any) => void;
  */
 export type FilterDeletedEventCallback = (event: FilterDeletedEvent) => void;
 
+/**
+ * A type for connection callback functions (blockchain event handler).
+ */
+export type ConnectionCallback = (webSocket) => void;
+
 /**
  * A type for disconnection callback functions (blockchain event handler).
  */