From ede15cc0e856b058e78904f9315a3fd12666bcc4 Mon Sep 17 00:00:00 2001
From: Harunobu Ishii <harunobuishii547@gmail.com>
Date: Thu, 28 Mar 2024 16:03:00 -0400
Subject: [PATCH] Method for Socket disconnection added

---
 lib/jiff-client.js                    |  4 ++++
 tests/arithmetics/arithmetics.test.ts | 21 +++++----------------
 2 files changed, 9 insertions(+), 16 deletions(-)

diff --git a/lib/jiff-client.js b/lib/jiff-client.js
index 47121b2c7..f2647f893 100644
--- a/lib/jiff-client.js
+++ b/lib/jiff-client.js
@@ -464,6 +464,10 @@ function JIFFClient(hostname, computation_id, options) {
     }
   };
 
+  this.disconnect = async function(){
+    await this.socket.safe_disconnect(true)
+  }
+
   // Connect when all is done
   if (options.autoConnect !== false) {
     this.connect();
diff --git a/tests/arithmetics/arithmetics.test.ts b/tests/arithmetics/arithmetics.test.ts
index b3f41ece3..493fc4929 100644
--- a/tests/arithmetics/arithmetics.test.ts
+++ b/tests/arithmetics/arithmetics.test.ts
@@ -48,23 +48,12 @@ describe('JIFF Arithmetic Operations', () => {
   });
 
   afterEach(async () => {
-    // Shutting Server
-    await jiffClient1.socket.onclose('io client disconnect')
-    await jiffClient2.socket.onclose('io client disconnect')
+    // Shutting down client
+    await jiffClient1.socket.disconnect()
+    await jiffClient2.socket.disconnect()
+    
+    // Shutting down Server
     await jiffServer.closeAllSockets();
-    // jiffServer.freeComputation(computation_id);
-  
-    // Close the server
-    // await new Promise((resolve, reject) => {
-    //   server.close((err:string) => {
-    //     if (err) {
-    //       console.error('Error closing server:', err);
-    //       reject(err);
-    //       return;
-    //     }
-    //     resolve(console.log('Server closed'));
-    //   });
-    // });
   });
 
   it('should correctly add 60 + 60 = 120', async () => {