1 files changed, 32 insertions, 29 deletions

diff --git a/test/node/extensions.js b/test/node/extensions.js
index 35b4ac9..48b1dfd 100644
--- a/test/node/extensions.js
+++ b/test/node/extensions.js
@@ -2,55 +2,58 @@ const fixtures = require('webtorrent-fixtures')
 const test = require('tape')
 const WebTorrent = require('../../')
 
-test('extension support', function (t) {
+test('extension support', t => {
   t.plan(6)
   let extendedHandshakes = 0
 
-  function Extension (wire) {
-    wire.extendedHandshake.test = 'Hello, World!'
-  }
+  class Extension {
+    constructor (wire) {
+      wire.extendedHandshake.test = 'Hello, World!'
+    }
 
-  Extension.prototype.name = 'wt_test'
-  Extension.prototype.onExtendedHandshake = function (extendedHandshake) {
-    extendedHandshakes += 1
-
-    t.equal(
-      extendedHandshake.test.toString(), 'Hello, World!',
-      'handshake.test === Hello, World!'
-    )
-
-    if (extendedHandshakes === 2) {
-      client1.destroy(function (err) {
-        t.error(err, 'client1 destroyed')
-      })
-      client2.destroy(function (err) {
-        t.error(err, 'client2 destroyed')
-      })
+    onExtendedHandshake (extendedHandshake) {
+      extendedHandshakes += 1
+
+      t.equal(
+        extendedHandshake.test.toString(), 'Hello, World!',
+        'handshake.test === Hello, World!'
+      )
+
+      if (extendedHandshakes === 2) {
+        client1.destroy(err => {
+          t.error(err, 'client1 destroyed')
+        })
+        client2.destroy(err => {
+          t.error(err, 'client2 destroyed')
+        })
+      }
     }
   }
 
+  Extension.prototype.name = 'wt_test'
+
   const client1 = new WebTorrent({ dht: false, tracker: false, lsd: false })
 
-  client1.on('error', function (err) { t.fail(err) })
-  client1.on('warning', function (err) { t.fail(err) })
+  client1.on('error', err => { t.fail(err) })
+  client1.on('warning', err => { t.fail(err) })
 
   const client2 = new WebTorrent({ dht: false, tracker: false, lsd: false })
 
-  client2.on('error', function (err) { t.fail(err) })
-  client2.on('warning', function (err) { t.fail(err) })
+  client2.on('error', err => { t.fail(err) })
+  client2.on('warning', err => { t.fail(err) })
 
-  client1.add(fixtures.leaves.parsedTorrent, function (torrent1) {
-    torrent1.on('wire', function (wire) {
+  client1.add(fixtures.leaves.parsedTorrent, torrent1 => {
+    torrent1.on('wire', wire => {
       t.pass('client1 onWire')
       wire.use(Extension)
     })
     const torrent2 = client2.add(fixtures.leaves.parsedTorrent.infoHash)
-    torrent2.on('wire', function (wire) {
+    torrent2.on('wire', wire => {
       t.pass('client2 onWire')
       wire.use(Extension)
     })
-    torrent2.on('infoHash', function () {
-      torrent2.addPeer('127.0.0.1:' + client1.address().port)
+    torrent2.on('infoHash', () => {
+      torrent2.addPeer(`127.0.0.1:${client1.address().port}`)
     })
   })
 })