Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
ulixee
GitHub Repository: ulixee/secret-agent
 Path: blob/main/full-client/test/websocket.test.ts
1028 views
1
import { Helpers } from '@secret-agent/testing';

2
import MitmServer from '@secret-agent/mitm/lib/MitmProxy';

3
import { createPromise } from '@secret-agent/commons/utils';

4
import * as WebSocket from 'ws';

5
import HttpUpgradeHandler from '@secret-agent/mitm/handlers/HttpUpgradeHandler';

6
import WebsocketResource from '@secret-agent/client/lib/WebsocketResource';

7
import { ITestKoaServer } from '@secret-agent/testing/helpers';

8
import { AddressInfo } from 'net';

9
import Core from '@secret-agent/core/index';

10
import { Handler } from '../index';

11


12
let handler: Handler;

13
let koaServer: ITestKoaServer;

14
beforeAll(async () => {

15
  await Core.start();

16
  handler = new Handler({ host: await Core.server.address });

17
  Helpers.onClose(() => handler.close(), true);

18
  koaServer = await Helpers.runKoaServer();

19
});

20


21
afterAll(Helpers.afterAll);

22
afterEach(Helpers.afterEach);

23


24
describe('Websocket tests', () => {

25
  it('can wait for a websocket', async () => {

26
    const mitmServer = await MitmServer.start();

27
    const upgradeSpy = jest.spyOn(HttpUpgradeHandler.prototype, 'onUpgrade');

28
    Helpers.needsClosing.push(mitmServer);

29


30
    const serverMessagePromise = createPromise();

31
    const wss = new WebSocket.Server({ noServer: true });

32


33
    const receivedMessages: string[] = [];

34
    koaServer.server.on('upgrade', (request, socket, head) => {

35
      wss.handleUpgrade(request, socket, head, async (ws: WebSocket) => {

36
        ws.on('message', msg => {

37
          receivedMessages.push(msg.toString());

38
          if (msg === 'Echo Message19') {

39
            serverMessagePromise.resolve();

40
          }

41
        });

42
        for (let i = 0; i < 20; i += 1) {

43
          ws.send(`Message${i}`);

44
          await new Promise(setImmediate);

45
        }

46
      });

47
    });

48


49
    koaServer.get('/ws-test', async ctx => {

50
      ctx.body = `

51
<html lang="en">

52
  <body>

53
  <h1>Here we go</h1>

54
  </body>

55
  <script>

56
    const ws = new WebSocket('ws://localhost:${(koaServer.server.address() as AddressInfo).port}');

57
    ws.onmessage = msg => {

58
      ws.send('Echo ' + msg.data);

59
    };

60
    let hasRun = false;

61
    document.addEventListener('mousemove', () => {

62
      if (hasRun) return;

63
      hasRun = true;

64
      ws.send('Final message');

65
    })

66
  </script>

67
</html>`;

68
    });

69
    const agent = await handler.createAgent();

70


71
    await agent.goto(`${koaServer.baseUrl}/ws-test`);

72


73
    await agent.waitForElement(agent.document.querySelector('h1'));

74
    await serverMessagePromise.promise;

75
    expect(receivedMessages).toHaveLength(20);

76


77
    expect(upgradeSpy).toHaveBeenCalledTimes(1);

78


79
    const resources = await agent.waitForResource({ type: 'Websocket' });

80
    expect(resources).toHaveLength(1);

81


82
    const [wsResource] = resources as WebsocketResource[];

83


84
    const broadcast = createPromise();

85
    let messagesCtr = 0;

86
    await wsResource.on('message', message => {

87
      messagesCtr += 1;

88
      if (message.message === 'Final message') {

89
        broadcast.resolve();

90
      }

91
    });

92
    await agent.interact({ move: [10, 10] });

93
    await broadcast.promise;

94
    expect(messagesCtr).toBe(41);

95


96
    await agent.close();

97
  });

98
});

99


100