diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 3aedb75..c6da9a9 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -27,4 +27,7 @@ jobs:
           make build
 
       - name: Run linter
-        run: make lint
\ No newline at end of file
+        run: make lint
+
+      - name: Run test
+        run: make test
\ No newline at end of file
diff --git a/__fixtures__/getRemoveResponse.js b/__fixtures__/getRemoveResponse.js
index 8667533..9d13b1b 100644
--- a/__fixtures__/getRemoveResponse.js
+++ b/__fixtures__/getRemoveResponse.js
@@ -1,5 +1,3 @@
-const getRemoveResponse = gameId => ({
-  data: JSON.stringify([null, null, 'lobby', 'game:remove', { id: gameId }]),
-});
+const getRemoveResponse = gameId => (JSON.stringify([null, null, 'lobby', 'game:remove', { id: gameId }]));
 
 export default getRemoveResponse;
diff --git a/__fixtures__/getUpdateResponse.js b/__fixtures__/getUpdateResponse.js
index 7eaa585..4663f2c 100644
--- a/__fixtures__/getUpdateResponse.js
+++ b/__fixtures__/getUpdateResponse.js
@@ -1,36 +1,34 @@
-const getUpdateResponse = (gameId, state = 'waiting_opponent') => ({
-  data: JSON.stringify([null,
-    null,
-    'lobby',
-    'game:upsert',
-    {
-      game: {
-        id: gameId,
-        inserted_at: '2020-10-04T22:38:22',
+const getUpdateResponse = (gameId, state = 'waiting_opponent') => (JSON.stringify([null,
+  null,
+  'lobby',
+  'game:upsert',
+  {
+    game: {
+      id: gameId,
+      inserted_at: '2020-10-04T22:38:22',
+      is_bot: false,
+      level: 'elementary',
+      players: [{
+        id: 2782,
+        name: 'H9ko',
         is_bot: false,
-        level: 'elementary',
-        players: [{
-          id: 2782,
-          name: 'H9ko',
-          is_bot: false,
-          github_id: 57991929,
-          lang: 'js',
-          editor_text: 'module.exports = () => {\n\n};',
-          editor_lang: 'js',
-          creator: true,
-          game_result: 'undefined',
-          check_result: {
-            asserts: [], asserts_count: 0, output: '', result: '{"status": "info"}', status: 'initial', success_count: 0,
-          },
-          achievements: ['played_ten_games', 'played_fifty_games'],
-          rating: 1152,
-          rating_diff: 0,
-        }],
-        state,
-        timeout_seconds: 7200,
-        type: 'public',
-      },
-    }]),
-});
+        github_id: 57991929,
+        lang: 'js',
+        editor_text: 'module.exports = () => {\n\n};',
+        editor_lang: 'js',
+        creator: true,
+        game_result: 'undefined',
+        check_result: {
+          asserts: [], asserts_count: 0, output: '', result: '{"status": "info"}', status: 'initial', success_count: 0,
+        },
+        achievements: ['played_ten_games', 'played_fifty_games'],
+        rating: 1152,
+        rating_diff: 0,
+      }],
+      state,
+      timeout_seconds: 7200,
+      type: 'public',
+    },
+  }]));
 
 export default getUpdateResponse;
diff --git a/__fixtures__/phxReply4bots.js b/__fixtures__/phxReply4bots.js
index 031b81e..21b0719 100644
--- a/__fixtures__/phxReply4bots.js
+++ b/__fixtures__/phxReply4bots.js
@@ -1,159 +1,157 @@
-export default {
-  data: JSON.stringify([
-    '7',
-    '7',
-    'lobby',
-    'phx_reply',
-    {
-      response: {
-        active_games: [
-          {
-            id: 19887,
-            inserted_at: '2020-09-27T13:30:54',
-            is_bot: true,
-            level: 'medium',
-            players: [
-              {
-                id: -27,
-                name: 'HappyQleaner',
-                is_bot: true,
-                github_id: 35539033,
-                lang: 'js',
-                editor_text: 'module.exports = () => {\n\n};',
-                editor_lang: 'js',
-                creator: true,
-                game_result: 'undefined',
-                check_result: {
-                  asserts: [],
-                  asserts_count: 0,
-                  output: '',
-                  result: '{"status": "info"}',
-                  status: 'initial',
-                  success_count: 0,
-                },
-                achievements: [
-                  'played_hundred_games',
-                  'win_games_with?clojure_cpp_csharp_golang_haskell_js_kotlin_php_python_ruby',
-                ],
-                rating: 1365,
-                rating_diff: 0,
+export default JSON.stringify([
+  '7',
+  '7',
+  'lobby',
+  'phx_reply',
+  {
+    response: {
+      active_games: [
+        {
+          id: 19887,
+          inserted_at: '2020-09-27T13:30:54',
+          is_bot: true,
+          level: 'medium',
+          players: [
+            {
+              id: -27,
+              name: 'HappyQleaner',
+              is_bot: true,
+              github_id: 35539033,
+              lang: 'js',
+              editor_text: 'module.exports = () => {\n\n};',
+              editor_lang: 'js',
+              creator: true,
+              game_result: 'undefined',
+              check_result: {
+                asserts: [],
+                asserts_count: 0,
+                output: '',
+                result: '{"status": "info"}',
+                status: 'initial',
+                success_count: 0,
               },
-            ],
-            state: 'waiting_opponent',
-            timeout_seconds: 3600,
-            type: 'bot',
-          },
-          {
-            id: 19923,
-            inserted_at: '2020-09-28T12:17:11',
-            is_bot: true,
-            level: 'hard',
-            players: [
-              {
-                id: -18,
-                name: 'DenisPython',
-                is_bot: true,
-                github_id: 35539033,
-                lang: 'js',
-                editor_text: 'module.exports = () => {\n\n};',
-                editor_lang: 'js',
-                creator: true,
-                game_result: 'undefined',
-                check_result: {
-                  asserts: [],
-                  asserts_count: 0,
-                  output: '',
-                  result: '{"status": "info"}',
-                  status: 'initial',
-                  success_count: 0,
-                },
-                achievements: [
-                  'played_hundred_games',
-                  'win_games_with?cpp_csharp_elixir_golang_haskell_java_js_php_python_ruby_ts',
-                ],
-                rating: 1344,
-                rating_diff: 0,
+              achievements: [
+                'played_hundred_games',
+                'win_games_with?clojure_cpp_csharp_golang_haskell_js_kotlin_php_python_ruby',
+              ],
+              rating: 1365,
+              rating_diff: 0,
+            },
+          ],
+          state: 'waiting_opponent',
+          timeout_seconds: 3600,
+          type: 'bot',
+        },
+        {
+          id: 19923,
+          inserted_at: '2020-09-28T12:17:11',
+          is_bot: true,
+          level: 'hard',
+          players: [
+            {
+              id: -18,
+              name: 'DenisPython',
+              is_bot: true,
+              github_id: 35539033,
+              lang: 'js',
+              editor_text: 'module.exports = () => {\n\n};',
+              editor_lang: 'js',
+              creator: true,
+              game_result: 'undefined',
+              check_result: {
+                asserts: [],
+                asserts_count: 0,
+                output: '',
+                result: '{"status": "info"}',
+                status: 'initial',
+                success_count: 0,
               },
-            ],
-            state: 'waiting_opponent',
-            timeout_seconds: 3600,
-            type: 'bot',
-          },
-          {
-            id: 20011,
-            inserted_at: '2020-09-30T16:15:49',
-            is_bot: true,
-            level: 'elementary',
-            players: [
-              {
-                id: -9,
-                name: 'UlaBack',
-                is_bot: true,
-                github_id: 35539033,
-                lang: 'ruby',
-                editor_text: 'module.exports = () => {\n\n};',
-                editor_lang: 'ruby',
-                creator: true,
-                game_result: 'undefined',
-                check_result: {
-                  asserts: [],
-                  asserts_count: 0,
-                  output: '',
-                  result: '{"status": "info"}',
-                  status: 'initial',
-                  success_count: 0,
-                },
-                achievements: [
-                  'played_hundred_games',
-                  'win_games_with?cpp_csharp_elixir_haskell_java_js_php_python_ruby',
-                ],
-                rating: 1261,
-                rating_diff: 0,
+              achievements: [
+                'played_hundred_games',
+                'win_games_with?cpp_csharp_elixir_golang_haskell_java_js_php_python_ruby_ts',
+              ],
+              rating: 1344,
+              rating_diff: 0,
+            },
+          ],
+          state: 'waiting_opponent',
+          timeout_seconds: 3600,
+          type: 'bot',
+        },
+        {
+          id: 20011,
+          inserted_at: '2020-09-30T16:15:49',
+          is_bot: true,
+          level: 'elementary',
+          players: [
+            {
+              id: -9,
+              name: 'UlaBack',
+              is_bot: true,
+              github_id: 35539033,
+              lang: 'ruby',
+              editor_text: 'module.exports = () => {\n\n};',
+              editor_lang: 'ruby',
+              creator: true,
+              game_result: 'undefined',
+              check_result: {
+                asserts: [],
+                asserts_count: 0,
+                output: '',
+                result: '{"status": "info"}',
+                status: 'initial',
+                success_count: 0,
               },
-            ],
-            state: 'waiting_opponent',
-            timeout_seconds: 3600,
-            type: 'bot',
-          },
-          {
-            id: 20013,
-            inserted_at: '2020-09-30T17:14:20',
-            is_bot: true,
-            level: 'easy',
-            players: [
-              {
-                id: -5,
-                name: 'ValyaFront',
-                is_bot: true,
-                github_id: 35539033,
-                lang: 'js',
-                editor_text: 'module.exports = () => {\n\n};',
-                editor_lang: 'js',
-                creator: true,
-                game_result: 'undefined',
-                check_result: {
-                  asserts: [],
-                  asserts_count: 0,
-                  output: '',
-                  result: '{"status": "info"}',
-                  status: 'initial',
-                  success_count: 0,
-                },
-                achievements: [
-                  'played_hundred_games',
-                  'win_games_with?clojure_cpp_elixir_golang_java_js_php_python_ruby',
-                ],
-                rating: 1287,
-                rating_diff: 0,
+              achievements: [
+                'played_hundred_games',
+                'win_games_with?cpp_csharp_elixir_haskell_java_js_php_python_ruby',
+              ],
+              rating: 1261,
+              rating_diff: 0,
+            },
+          ],
+          state: 'waiting_opponent',
+          timeout_seconds: 3600,
+          type: 'bot',
+        },
+        {
+          id: 20013,
+          inserted_at: '2020-09-30T17:14:20',
+          is_bot: true,
+          level: 'easy',
+          players: [
+            {
+              id: -5,
+              name: 'ValyaFront',
+              is_bot: true,
+              github_id: 35539033,
+              lang: 'js',
+              editor_text: 'module.exports = () => {\n\n};',
+              editor_lang: 'js',
+              creator: true,
+              game_result: 'undefined',
+              check_result: {
+                asserts: [],
+                asserts_count: 0,
+                output: '',
+                result: '{"status": "info"}',
+                status: 'initial',
+                success_count: 0,
               },
-            ],
-            state: 'waiting_opponent',
-            timeout_seconds: 3600,
-            type: 'bot',
-          },
-        ],
-      },
-      status: 'ok',
+              achievements: [
+                'played_hundred_games',
+                'win_games_with?clojure_cpp_elixir_golang_java_js_php_python_ruby',
+              ],
+              rating: 1287,
+              rating_diff: 0,
+            },
+          ],
+          state: 'waiting_opponent',
+          timeout_seconds: 3600,
+          type: 'bot',
+        },
+      ],
     },
-  ]),
-};
+    status: 'ok',
+  },
+]);
diff --git a/__test__/background.test.js b/__test__/background.test.js
index dfd69cb..31e7d40 100644
--- a/__test__/background.test.js
+++ b/__test__/background.test.js
@@ -1,46 +1,74 @@
 import browser from 'webextension-polyfill';
+import WS from 'jest-websocket-mock';
 import socketConnect from '../src/background/socket';
 import gameStatuses from '../src/background/models';
 import phxReply4bots from '../__fixtures__/phxReply4bots';
 import getUpdateResponseWithID from '../__fixtures__/getUpdateResponse';
 import getRemoveResponseWithID from '../__fixtures__/getRemoveResponse';
 
-describe('socket', () => {
-  const fakeSocket = {};
-  socketConnect(('wss://codebattle.hexlet.io/extension/websocket?vsn=2.0.0'), fakeSocket);
+let serverWS = null;
+beforeEach(async () => {
+  jest.clearAllMocks();
+  serverWS = new WS('ws://localhost:1234');
+  serverWS.on('connection', socket => {
+    socket.on('message', data => {
+      switch (data) {
+        case ['7', '7', 'lobby', 'phx_join', {}]:
+          serverWS.send(phxReply4bots);
+          break;
+        case [null, '8', 'phoenix', 'heartbeat', {}]:
+          serverWS.send([null, '8', 'phoenix', 'heartbeat', {}]);
+          break;
+        default:
+          break;
+      }
+    });
+  });
+  socketConnect('ws://localhost:1234');
+  await serverWS.connected;
+});
 
+afterEach(() => {
+  WS.clean();
+});
+
+describe('socket', () => {
   test('phxReply', () => {
-    fakeSocket.onmessage(phxReply4bots);
+    expect(browser.browserAction.setBadgeText).toHaveBeenCalledTimes(0);
+    serverWS.send(phxReply4bots);
     expect(browser.browserAction.setBadgeText).toHaveBeenLastCalledWith({
       text: null,
     });
   });
   test('add games->remove games', () => {
-    fakeSocket.onmessage(getUpdateResponseWithID(1000, gameStatuses.waiting));
+    expect(browser.browserAction.setBadgeText).toHaveBeenCalledTimes(0);
+    serverWS.send(getUpdateResponseWithID(1000, gameStatuses.waiting));
     expect(browser.browserAction.setBadgeText).toHaveBeenLastCalledWith({
       text: '1',
     });
-    fakeSocket.onmessage(getUpdateResponseWithID(1001, gameStatuses.waiting));
+    serverWS.send(getUpdateResponseWithID(1001, gameStatuses.waiting));
     expect(browser.browserAction.setBadgeText).toHaveBeenLastCalledWith({
       text: '2',
     });
-    fakeSocket.onmessage(getRemoveResponseWithID(1001));
+    serverWS.send(getRemoveResponseWithID(1001));
     expect(browser.browserAction.setBadgeText).toHaveBeenLastCalledWith({
       text: '1',
     });
-    fakeSocket.onmessage(getRemoveResponseWithID(1000));
+    serverWS.send(getRemoveResponseWithID(1000));
     expect(browser.browserAction.setBadgeText).toHaveBeenLastCalledWith({
       text: null,
     });
   });
   test('add game -> changeStatus', () => {
-    fakeSocket.onmessage(getUpdateResponseWithID(1002, gameStatuses.waiting));
+    expect(browser.browserAction.setBadgeText).toHaveBeenCalledTimes(0);
+    serverWS.send(getUpdateResponseWithID(1002, gameStatuses.waiting));
     expect(browser.browserAction.setBadgeText).toHaveBeenLastCalledWith({
       text: '1',
     });
-    fakeSocket.onmessage(getUpdateResponseWithID(1002, gameStatuses.joined));
+    serverWS.send(getUpdateResponseWithID(1002, gameStatuses.joined));
     expect(browser.browserAction.setBadgeText).toHaveBeenLastCalledWith({
       text: null,
     });
+    expect(browser.browserAction.setBadgeText).toHaveBeenCalledTimes(2);
   });
 });
diff --git a/babel.config.js b/babel.config.js
index f037a1a..9c67a0d 100644
--- a/babel.config.js
+++ b/babel.config.js
@@ -1,12 +1,7 @@
 module.exports = {
-  presets: [
-    [
-      '@babel/preset-env',
-      {
-        targets: {
-          node: 'current',
-        },
-      },
-    ],
-  ],
+  env: {
+    test: {
+      plugins: ['@babel/plugin-transform-runtime'],
+    },
+  },
 };
diff --git a/package-lock.json b/package-lock.json
index 508a0b6..996eb86 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -9213,6 +9213,12 @@
       "integrity": "sha512-pvJFFPr2El9inccdFeNbGXVDhaTytzOF/1S049NtEDD1pMCkreeXYWLhM2Xca6zudge4GApzfrCkrxG986MiHQ==",
       "dev": true
     },
+    "jest-websocket-mock": {
+      "version": "2.2.0",
+      "resolved": "https://registry.npmjs.org/jest-websocket-mock/-/jest-websocket-mock-2.2.0.tgz",
+      "integrity": "sha512-lc3wwXOEyNa4ZpcgJtUG3mmKMAq5FAsKYiZph0p/+PAJrAPuX4JCIfJMdJ/urRsLBG51fwm/wlVPNbR6s2nzNw==",
+      "dev": true
+    },
     "jest-worker": {
       "version": "26.3.0",
       "resolved": "https://registry.npmjs.org/jest-worker/-/jest-worker-26.3.0.tgz",
@@ -9957,6 +9963,15 @@
         "minimist": "^1.2.5"
       }
     },
+    "mock-socket": {
+      "version": "9.0.3",
+      "resolved": "https://registry.npmjs.org/mock-socket/-/mock-socket-9.0.3.tgz",
+      "integrity": "sha512-SxIiD2yE/By79p3cNAAXyLQWTvEFNEzcAO7PH+DzRqKSFaplAPFjiQLmw8ofmpCsZf+Rhfn2/xCJagpdGmYdTw==",
+      "dev": true,
+      "requires": {
+        "url-parse": "^1.4.4"
+      }
+    },
     "move-concurrently": {
       "version": "1.0.1",
       "resolved": "https://registry.npmjs.org/move-concurrently/-/move-concurrently-1.0.1.tgz",
@@ -11312,6 +11327,12 @@
       "integrity": "sha1-nsYfeQSYdXB9aUFFlv2Qek1xHnM=",
       "dev": true
     },
+    "querystringify": {
+      "version": "2.2.0",
+      "resolved": "https://registry.npmjs.org/querystringify/-/querystringify-2.2.0.tgz",
+      "integrity": "sha512-FIqgj2EUvTa7R50u0rGsyTftzjYmv/a3hO345bZNrqabNqjtgiDMgmo4mkUjd+nzU5oF3dClKqFIPUKybUyqoQ==",
+      "dev": true
+    },
     "randombytes": {
       "version": "2.1.0",
       "resolved": "https://registry.npmjs.org/randombytes/-/randombytes-2.1.0.tgz",
@@ -11702,6 +11723,12 @@
       "integrity": "sha512-NKN5kMDylKuldxYLSUfrbo5Tuzh4hd+2E8NPPX02mZtn1VuREQToYe/ZdlJy+J3uCpfaiGF05e7B8W0iXbQHmg==",
       "dev": true
     },
+    "requires-port": {
+      "version": "1.0.0",
+      "resolved": "https://registry.npmjs.org/requires-port/-/requires-port-1.0.0.tgz",
+      "integrity": "sha1-kl0mAdOaxIXgkc8NpcbmlNw9yv8=",
+      "dev": true
+    },
     "resolve": {
       "version": "1.17.0",
       "resolved": "https://registry.npmjs.org/resolve/-/resolve-1.17.0.tgz",
@@ -13530,6 +13557,16 @@
         }
       }
     },
+    "url-parse": {
+      "version": "1.4.7",
+      "resolved": "https://registry.npmjs.org/url-parse/-/url-parse-1.4.7.tgz",
+      "integrity": "sha512-d3uaVyzDB9tQoSXFvuSUNFibTd9zxd2bkVrDRvF5TmvWWQwqE4lgYJ5m+x1DbecWkw+LK4RNl2CU1hHuOKPVlg==",
+      "dev": true,
+      "requires": {
+        "querystringify": "^2.1.1",
+        "requires-port": "^1.0.0"
+      }
+    },
     "use": {
       "version": "3.1.1",
       "resolved": "https://registry.npmjs.org/use/-/use-3.1.1.tgz",
diff --git a/package.json b/package.json
index 9bc636d..a30611a 100644
--- a/package.json
+++ b/package.json
@@ -58,7 +58,9 @@
     "html-webpack-plugin": "^4.3.0",
     "jest": "^26.4.2",
     "jest-webextension-mock": "^3.6.1",
+    "jest-websocket-mock": "^2.2.0",
     "mini-css-extract-plugin": "^0.9.0",
+    "mock-socket": "^9.0.3",
     "node-sass": "^4.14.1",
     "sass-loader": "^8.0.2",
     "webpack": "^4.43.0",