From fbe3b0b3f6016b719b13026d772bfc485a3dc10a Mon Sep 17 00:00:00 2001 From: Keaton Sentak Date: Tue, 8 Aug 2023 15:52:04 -0400 Subject: [PATCH] move ws and filestream into one class --- server/src/sessionManagement.mjs | 116 ++++---- server/test/suite/sessionManagement.test.mjs | 266 ++++++++++++------- 2 files changed, 212 insertions(+), 170 deletions(-) diff --git a/server/src/sessionManagement.mjs b/server/src/sessionManagement.mjs index b12dd328..a751b01c 100644 --- a/server/src/sessionManagement.mjs +++ b/server/src/sessionManagement.mjs @@ -25,56 +25,61 @@ import fs from 'fs'; import yaml from 'js-yaml'; import WebSocket from 'ws'; -class SessionWebSocket { +class SessionHandler { constructor() { + this.mode = null; this.ws = null; + this.stream = null; } - open(dir) { - this.ws = new WebSocket(dir); - this.ws.on('open', () => { - logger.info(`Websocket connection opened: ${dir}`); - }) - } - - close() { - if (this.ws) { - this.ws.close(); - this.ws = null; + // Determine mode based on directory/url + _determineMode(dir) { + const wsRegex = /^(ws(s)?):\/\//i; + if (wsRegex.test(dir)) { + this.mode = 'websocket'; + } else { + this.mode = 'filestream'; } } -} - -class SessionFileStream { - constructor() { - this.stream = null; - } open(dir) { - if (!fs.existsSync(dir)) { - logger.info("Directory does not exist for: " + dir); - fs.mkdirSync(dir, { recursive: true}); + this._determineMode(dir); + + if (this.mode === 'websocket') { + this.ws = new WebSocket(dir); + this.ws.on('open', () => { + logger.info(`Websocket connection opened: ${dir}`); + }); + } else { + if (!fs.existsSync(dir)) { + logger.info("Directory does not exist for: " + dir); + fs.mkdirSync(dir, { recursive: true }); + } + + this.stream = fs.createWriteStream(`${dir}/FireboltCalls_live.log`, { flags: 'a' }); } - - this.stream = fs.createWriteStream(`${dir}/FireboltCalls_live.log`, { flags: 'a' }); } write(data) { - if (this.stream) { + if (this.mode === 'websocket' && this.ws) { + this.ws.send(data); + } else if (this.stream) { this.stream.write(`${data}\n`); } } close() { - if (this.stream) { + if (this.mode === 'websocket' && this.ws) { + this.ws.close(); + this.ws = null; + } else if (this.stream) { this.stream.end(); this.stream = null; } } } -let sessionWebSocket = new SessionWebSocket(); -let sessionFileStream = new SessionFileStream(); +let sessionHandler = new SessionHandler(); class FireboltCall { constructor(methodCall, params) { @@ -111,6 +116,10 @@ class Session { // const sessionStartString = sessionStart.toISOString().replace(/T/, '_').replace(/\..+/, ''); // logger.info(`${sessionStart.toISOString()}`); + // Check if the output path is a WebSocket URL + const wsRegex = /^(ws(s)?):\/\//i; + this.sessionOutputPath = wsRegex.test(this.sessionOutputPath) ? "./sessions/output" : this.sessionOutputPath; + if (!fs.existsSync(this.sessionOutputPath)) { logger.info("Directory does not exist for: " + this.sessionOutputPath) fs.mkdirSync(this.sessionOutputPath, { recursive: true}); @@ -408,12 +417,7 @@ function stopRecording() { logger.info('Stopping recording'); sessionRecording.recording = false; const sessionData = sessionRecording.recordedSession.exportSession(); - if (sessionWebSocket.ws) { - sessionWebSocket.close(); - } - if (sessionFileStream.stream) { - sessionFileStream.close(); - } + sessionHandler.close(); return sessionData; } else { logger.warn('Trying to stop recording when not recording'); @@ -432,11 +436,7 @@ function addCall(methodCall, params){ sessionRecording.recordedSession.calls.push(call); if (sessionRecording.recordedSession.sessionOutput === "live") { const data = JSON.stringify(call); - if (sessionWebSocket.ws) { - sessionWebSocket.ws.send(data); - } else if (sessionFileStream.stream) { - sessionFileStream.write(data); - } + sessionHandler.write(data); } } } @@ -451,18 +451,12 @@ function getOutputFormat(){ } function setOutputDir(dir) { - const wsRegex = /^(ws(s)?):\/\//i; - - if (wsRegex.test(dir)) { - sessionWebSocket.open(dir); - } else { - sessionRecording.recordedSession.sessionOutputPath = dir; - sessionRecording.recordedSession.mockOutputPath = dir; - logger.info("Setting output path: " + sessionRecording.recordedSession.mockOutputPath); - if (sessionRecording.recordedSession.sessionOutput === "live") { - sessionFileStream.open(dir); - } - } + if (sessionRecording.recordedSession.sessionOutput === "live") { + sessionHandler.open(dir); + } + sessionRecording.recordedSession.sessionOutputPath = dir; + sessionRecording.recordedSession.mockOutputPath = dir; + logger.info("Setting output path: " + sessionRecording.recordedSession.mockOutputPath); } function getSessionOutputDir(){ @@ -482,11 +476,7 @@ function updateCallWithResponse(method, result, key) { sessionRecording.recordedSession.calls.concat(...methodCalls); if (sessionRecording.recordedSession.sessionOutput === "live") { const data = JSON.stringify(methodCalls[i]); - if (sessionWebSocket.ws) { - sessionWebSocket.ws.send(data); - } else if (sessionFileStream.stream) { - sessionFileStream.write(data); - } + sessionHandler.write(data); } } } @@ -494,24 +484,14 @@ function updateCallWithResponse(method, result, key) { } // Utility function for unit tests -const setTestEntity = (entityName, mockEntity) => { - switch (entityName) { - case 'websocket': - sessionWebSocket = mockEntity; - break; - case 'filestream': - sessionFileStream = mockEntity; - break; - default: - throw new Error('Unknown entity name'); - } +const setTestEntity = (mockEntity) => { + sessionHandler = mockEntity } export const testExports = { setTestEntity, setOutputDir, - SessionFileStream, - SessionWebSocket + SessionHandler } export { Session, FireboltCall, startRecording, setOutputDir, stopRecording, addCall, isRecording, updateCallWithResponse, setOutputFormat, getOutputFormat, getSessionOutputDir, getMockOutputDir }; \ No newline at end of file diff --git a/server/test/suite/sessionManagement.test.mjs b/server/test/suite/sessionManagement.test.mjs index 178cacbc..2b3ca2ce 100644 --- a/server/test/suite/sessionManagement.test.mjs +++ b/server/test/suite/sessionManagement.test.mjs @@ -133,90 +133,127 @@ describe(`FireboltCall`, () => { }); }); -describe('SessionWebSocket', () => { +describe(`SessionHandler`, () => { afterEach(() => { - jest.restoreAllMocks(); + jest.restoreAllMocks(); }); - test('should initialize without a WebSocket', () => { - const sessionWebSocket = new sessionManagement.testExports.SessionWebSocket(); - expect(sessionWebSocket.ws).toBeNull(); + test(`should initialize without a WebSocket or FileStream`, () => { + const sessionHandler = new sessionManagement.testExports.SessionHandler(); + expect(sessionHandler.ws).toBeNull(); + expect(sessionHandler.stream).toBeNull(); }); - test('should open a WebSocket connection', () => { - const mockWsInstance = { - open: jest.fn(), - close: jest.fn(), - ws: { - open: jest.fn(), - close: jest.fn() - } - }; - - const SessionWebSocket = jest.spyOn(sessionManagement.testExports, 'SessionWebSocket').mockImplementation(() => mockWsInstance); - const dir = 'ws://example.com'; - - const sessionWebSocket = new SessionWebSocket(); - sessionWebSocket.open(dir); - - expect(sessionWebSocket.ws).toBeTruthy(); - expect(sessionWebSocket.open).toHaveBeenCalled(); + describe(`Websocket functionality`, () => { + test(`should open a WebSocket connection`, () => { + const mockSessionHandler = { + ws: { + open: jest.fn(), + send: jest.fn(), + close: jest.fn(), + }, + stream: null, + close: jest.fn(() => { + mockSessionHandler.ws.close() + }), + write: jest.fn(() => { + mockSessionHandler.ws.send() + }), + open: jest.fn(() => { + mockSessionHandler.ws.open() + }), + mode: 'websocket' + }; + + const SessionHandler = jest.spyOn(sessionManagement.testExports, 'SessionHandler') + .mockImplementation(() => mockSessionHandler); + const dir = 'ws://example.com'; + + const sessionHandler = new SessionHandler(); + sessionHandler.open(dir); + + expect(sessionHandler.open).toHaveBeenCalled(); + expect(sessionHandler.ws.open).toHaveBeenCalled(); + }); + + test(`should close a WebSocket connection`, () => { + const mockSessionHandler = { + ws: { + open: jest.fn(), + send: jest.fn(), + close: jest.fn(), + }, + stream: null, + close: jest.fn(() => { + mockSessionHandler.ws = null + }), + write: jest.fn(() => { + mockSessionHandler.ws.send() + }), + open: jest.fn(() => { + mockSessionHandler.ws.open() + }), + mode: 'websocket' + }; + + const SessionHandler = jest.spyOn(sessionManagement.testExports, 'SessionHandler') + .mockImplementation(() => mockSessionHandler); + const dir = 'ws://example.com'; + + const sessionHandler = new SessionHandler(); + sessionHandler.close(dir); + + expect(sessionHandler.close).toHaveBeenCalled(); + expect(sessionHandler.ws).toBeNull(); + }); }); - test('should close a WebSocket connection', () => { - const mockWsInstance = { - open: jest.fn(), - close: jest.fn(() => { - mockWsInstance.ws = null; - }), - ws: { - open: jest.fn(), - close: jest.fn() - } - }; + describe(`FileStream functionality`, () => { + test(`should open a file stream`, () => { + const dir = './some/directory/path'; - const SessionWebSocket = jest.spyOn(sessionManagement.testExports, 'SessionWebSocket').mockImplementation(() => mockWsInstance); - const dir = 'ws://example.com'; - - const sessionWebSocket = new SessionWebSocket(); - sessionWebSocket.close(dir); + const sessionHandler = new sessionManagement.testExports.SessionHandler(); + sessionHandler.open(dir); - expect(sessionWebSocket.ws).toBeNull(); - }); -}); + expect(sessionHandler.stream).toBeTruthy(); + }); -describe(`SessionFileStream`, () => { - test('should initialize without a stream', () => { - const sessionFileStream = new sessionManagement.testExports.SessionFileStream(); - expect(sessionFileStream.stream).toBeNull(); - }); + test(`should write data to the file stream`, () => { + const dir = './some/directory/path'; + const data = 'test data'; - test('should open a file stream', () => { - const sessionFileStream = new sessionManagement.testExports.SessionFileStream(); - const dir = './some/directory/path'; + const mockWriteStream = { + write: jest.fn(), + end: jest.fn() + }; - sessionFileStream.open(dir); - - expect(sessionFileStream.stream).toBeTruthy(); - }); + jest.spyOn(fs, 'createWriteStream').mockReturnValue(mockWriteStream); - test('should write data to the file stream', () => { - const sessionFileStream = new sessionManagement.testExports.SessionFileStream(); - sessionFileStream.open('./some/directory/path'); + const sessionHandler = new sessionManagement.testExports.SessionHandler(); + sessionHandler.open(dir); + sessionHandler.write(data); - const data = 'test data'; - sessionFileStream.write(data); + expect(mockWriteStream.write).toHaveBeenCalledWith(`${data}\n`); + }); - expect(() => sessionFileStream.write(data)).not.toThrow(); - }); + test(`should close the file stream`, () => { + const dir = './some/directory/path'; + + const mockEnd = jest.fn(); + const mockWriteStream = { + write: jest.fn(), + end: mockEnd + }; - test('should close the file stream', () => { - const sessionFileStream = new sessionManagement.testExports.SessionFileStream(); - sessionFileStream.open('./some/directory/path'); + jest.spyOn(fs, 'createWriteStream').mockReturnValue(mockWriteStream); - sessionFileStream.close(); + const sessionHandler = new sessionManagement.testExports.SessionHandler(); + sessionHandler.open(dir); + sessionHandler.close(); - expect(sessionFileStream.stream).toBeNull(); + expect(mockEnd).toHaveBeenCalled(); + expect(sessionHandler.stream).toBeNull(); + }); }); }); @@ -243,28 +280,36 @@ test(`sessionManagement.stopRecording closes the WebSocket connection`, () => { .mockImplementation((dir) => { const wsRegex = /^(ws(s)?):\/\//i; if (wsRegex.test(dir)) { - mockSessionWebSocket.ws.open(); + mockSessionHandler.open(); } else { - mockSessionFileStream.open() + mockSessionHandler.open() } }); - const mockSessionWebSocket = { - ws: { - open: jest.fn(), - close: jest.fn() - }, - close: jest.fn() - }; + const mockSessionHandler = { + ws: { + open: jest.fn(), + send: jest.fn(), + close: jest.fn() + }, + close: jest.fn(() => { + mockSessionHandler.ws.close() + }), + write: jest.fn(() => { + mockSessionHandler.ws.send() + }), + open: jest.fn() + }; + - sessionManagement.testExports.setTestEntity('websocket', mockSessionWebSocket); + sessionManagement.testExports.setTestEntity(mockSessionHandler); sessionManagement.testExports.setOutputDir('ws://example.com'); sessionManagement.startRecording(); sessionManagement.stopRecording(); - expect(mockSessionWebSocket.close).toHaveBeenCalled(); + expect(mockSessionHandler.ws.close).toHaveBeenCalled(); spySetOutputDir.mockRestore(); }); @@ -274,26 +319,34 @@ test(`sessionManagement.stopRecording closes the FileStream connection`, () => { .mockImplementation((dir) => { const wsRegex = /^(ws(s)?):\/\//i; if (wsRegex.test(dir)) { - mockSessionWebSocket.ws.open(); + mockSessionHandler.open(); } else { - mockSessionFileStream.open() + mockSessionHandler.open() } }); - const mockSessionFileStream = { + const mockSessionHandler = { + stream: { + open: jest.fn(), + write: jest.fn(), + close: jest.fn(), + }, + close: jest.fn(() => { + mockSessionHandler.stream.close(); + }), + write: jest.fn(() => { + mockSessionHandler.stream.write(); + }), open: jest.fn(), - write: jest.fn(), - close: jest.fn(), - stream: {} }; - sessionManagement.testExports.setTestEntity('filestream', mockSessionFileStream); + sessionManagement.testExports.setTestEntity(mockSessionHandler); sessionManagement.testExports.setOutputDir('./some/directory/path'); sessionManagement.startRecording(); sessionManagement.stopRecording(); - expect(mockSessionFileStream.close).toHaveBeenCalled(); + expect(mockSessionHandler.stream.close).toHaveBeenCalled(); spySetOutputDir.mockRestore(); }); @@ -315,28 +368,32 @@ test(`sessionManagement.addCall works properly`, () => { expect(result).toBeUndefined(); }); -test(`sessionManagement.addCall calls sessionWebSocket.send`, () => { +test(`sessionManagement.addCall calls websocket`, () => { const spySetOutputDir = jest.spyOn(sessionManagement.testExports, 'setOutputDir') .mockImplementation((dir) => { const wsRegex = /^(ws(s)?):\/\//i; if (wsRegex.test(dir)) { - mockSessionWebSocket.ws.open(); + mockSessionHandler.open(); } else { - mockSessionFileStream.open() + mockSessionHandler.open() } }); - const mockSessionWebSocket = { + const mockSessionHandler = { ws: { open: jest.fn(), send: jest.fn(), close: jest.fn() }, - close: jest.fn() + close: jest.fn(), + write: jest.fn(() => { + mockSessionHandler.ws.send() + }), + open: jest.fn() }; - sessionManagement.testExports.setTestEntity('websocket', mockSessionWebSocket); + sessionManagement.testExports.setTestEntity(mockSessionHandler); sessionManagement.testExports.setOutputDir('ws://example.com'); sessionManagement.startRecording(); @@ -344,31 +401,36 @@ test(`sessionManagement.addCall calls sessionWebSocket.send`, () => { sessionManagement.addCall("methodName", "Parameters"); - expect(mockSessionWebSocket.ws.send).toHaveBeenCalled(); + expect(mockSessionHandler.ws.send).toHaveBeenCalled(); spySetOutputDir.mockRestore(); }); -test(`sessionManagement.addCall calls sessionFileStream.write`, () => { +test(`sessionManagement.addCall calls filestream`, () => { const spySetOutputDir = jest.spyOn(sessionManagement.testExports, 'setOutputDir') .mockImplementation((dir) => { const wsRegex = /^(ws(s)?):\/\//i; if (wsRegex.test(dir)) { - mockSessionWebSocket.ws.open(); + mockSessionHandler.open(); } else { - mockSessionFileStream.open() + mockSessionHandler.open() } }); - const mockSessionFileStream = { - open: jest.fn(), - write: jest.fn(), + const mockSessionHandler = { + stream: { + open: jest.fn(), + write: jest.fn(), + close: jest.fn() + }, close: jest.fn(), - stream: jest.fn() + write: jest.fn(() => { + mockSessionHandler.stream.write() + }), + open: jest.fn() }; - sessionManagement.testExports.setTestEntity('filestream', mockSessionFileStream); - sessionManagement.testExports.setTestEntity('websocket', {}); + sessionManagement.testExports.setTestEntity(mockSessionHandler); sessionManagement.testExports.setOutputDir('./test'); sessionManagement.startRecording(); @@ -376,7 +438,7 @@ test(`sessionManagement.addCall calls sessionFileStream.write`, () => { sessionManagement.addCall("methodName", "Parameters"); - expect(mockSessionFileStream.write).toHaveBeenCalled(); + expect(mockSessionHandler.stream.write).toHaveBeenCalled(); spySetOutputDir.mockRestore(); });