Add util to wrap window stubs

ts/test/session/crypto/MessageEncrypter_test.ts

@@ -1,11 +1,9 @@
 import { expect } from 'chai';
-import { ImportMock } from 'ts-mock-imports';
 import * as crypto from 'crypto';
 import * as sinon from 'sinon';
-import * as window from '../../../window';
 import { MessageEncrypter } from '../../../session/crypto';
 import { EncryptionType } from '../../../session/types/EncryptionType';
-import { Stubs } from '../../test-utils';
+import { Stubs, TestUtils } from '../../test-utils';
 import { UserUtil } from '../../../util';
 import { SignalService } from '../../../protobuf';
 
@@ -14,35 +12,35 @@ describe('MessageEncrypter', () => {
   const ourNumber = 'ourNumber';
 
   beforeEach(() => {
-    ImportMock.mockOther(window, 'libsignal', {
+    TestUtils.stubWindow('libsignal', {
       SignalProtocolAddress: sandbox.stub(),
       SessionCipher: Stubs.SessionCipherStub,
     } as any);
 
-    ImportMock.mockOther(window, 'textsecure', {
+    TestUtils.stubWindow('textsecure', {
       storage: {
         protocol: sandbox.stub(),
       },
     });
 
-    ImportMock.mockOther(window, 'Signal', {
+    TestUtils.stubWindow('Signal', {
       Metadata: {
         SecretSessionCipher: Stubs.SecretSessionCipherStub,
       },
     });
 
-    ImportMock.mockOther(window, 'libloki', {
+    TestUtils.stubWindow('libloki', {
       crypto: {
         FallBackSessionCipher: Stubs.FallBackSessionCipherStub,
       },
     });
 
-    ImportMock.mockFunction(UserUtil, 'getCurrentDevicePubKey', ourNumber);
+    sandbox.stub(UserUtil, 'getCurrentDevicePubKey').resolves(ourNumber);
   });
 
   afterEach(() => {
     sandbox.restore();
-    ImportMock.restore();
+    TestUtils.restoreStubs();
   });
 
   describe('EncryptionType', () => {

ts/test/test-utils/testUtils.ts

@@ -1,21 +1,42 @@
 import * as sinon from 'sinon';
 import { ImportMock } from 'ts-mock-imports';
-import * as Shape from '../../../js/modules/data';
+import * as DataShape from '../../../js/modules/data';
+import * as window from '../../window';
+
+const sandbox = sinon.createSandbox();
 
 // We have to do this in a weird way because Data uses module.exports
 //  which doesn't play well with sinon or ImportMock
 // tslint:disable-next-line: no-require-imports no-var-requires
 const Data = require('../../../js/modules/data');
-type DataFunction = typeof Shape;
+type DataFunction = typeof DataShape;
 
 /**
- * Mock a function inside Data.
+ * Stub a function inside Data.
  *
- * Note: This uses `ImportMock` so you will have to call `ImportMock.restore()` or `stub.restore()` after each test.
+ * Note: This uses a custom sandbox.
+ * Please call `restoreStubs()` or `stub.restore()` to restore original functionality.
  */
-export function mockData(
-  fn: keyof DataFunction,
-  returns?: any
-): sinon.SinonStub {
-  return ImportMock.mockFunction(Data, fn, returns);
+export function stubData(fn: keyof DataFunction): sinon.SinonStub {
+  return sandbox.stub(Data, fn);
+}
+
+type WindowFunction = typeof window;
+
+/**
+ * Stub a window object.
+ *
+ * Note: This uses a custom sandbox.
+ * Please call `restoreStubs()` or `stub.restore()` to restore original functionality.
+ */
+export function stubWindow<K extends keyof WindowFunction>(
+  fn: K,
+  replaceWith?: Partial<WindowFunction[K]>
+) {
+  return ImportMock.mockOther(window, fn, replaceWith);
+}
+
+export function restoreStubs() {
+  ImportMock.restore();
+  sandbox.restore();
 }