[converter] Fix graph model tests by using spyOnAllFunctions

tfjs-converter/src/executor/graph_model.ts

@@ -44,6 +44,7 @@ export class GraphModel implements InferenceModel {
   private initializer: GraphExecutor;
   private resourceManager: ResourceManager;
   private signature: tensorflow.ISignatureDef;
+  private readonly io: typeof io;
 
   // Returns the version information for the tensorflow model GraphDef.
   get modelVersion(): string {
@@ -89,7 +90,9 @@ export class GraphModel implements InferenceModel {
    */
   constructor(
       private modelUrl: string|io.IOHandler,
-      private loadOptions: io.LoadOptions = {}) {
+      private loadOptions: io.LoadOptions = {},
+      tfio = io) {
+    this.io = tfio;
     if (loadOptions == null) {
       this.loadOptions = {};
     }
@@ -102,13 +105,13 @@ export class GraphModel implements InferenceModel {
       // Path is an IO Handler.
       this.handler = path as io.IOHandler;
     } else if (this.loadOptions.requestInit != null) {
-      this.handler = io.browserHTTPRequest(path as string, this.loadOptions);
+      this.handler = this.io.browserHTTPRequest(path as string, this.loadOptions);
     } else {
-      const handlers = io.getLoadHandlers(path as string, this.loadOptions);
+      const handlers = this.io.getLoadHandlers(path as string, this.loadOptions);
       if (handlers.length === 0) {
         // For backward compatibility: if no load handler can be found,
         // assume it is a relative http path.
-        handlers.push(io.browserHTTPRequest(path as string, this.loadOptions));
+        handlers.push(this.io.browserHTTPRequest(path as string, this.loadOptions));
       } else if (handlers.length > 1) {
         throw new Error(
             `Found more than one (${handlers.length}) load handlers for ` +
@@ -157,7 +160,7 @@ export class GraphModel implements InferenceModel {
 
     this.version = `${graph.versions.producer}.${graph.versions.minConsumer}`;
     const weightMap =
-        io.decodeWeights(this.artifacts.weightData, this.artifacts.weightSpecs);
+        this.io.decodeWeights(this.artifacts.weightData, this.artifacts.weightSpecs);
     this.executor = new GraphExecutor(
         OperationMapper.Instance.transformGraph(graph, this.signature));
     this.executor.weightMap = this.convertTensorMapToTensorsMap(weightMap);
@@ -228,7 +231,7 @@ export class GraphModel implements InferenceModel {
   async save(handlerOrURL: io.IOHandler|string, config?: io.SaveConfig):
       Promise<io.SaveResult> {
     if (typeof handlerOrURL === 'string') {
-      const handlers = io.getSaveHandlers(handlerOrURL);
+      const handlers = this.io.getSaveHandlers(handlerOrURL);
       if (handlers.length === 0) {
         throw new Error(
             `Cannot find any save handlers for URL '${handlerOrURL}'`);
@@ -438,7 +441,7 @@ export class GraphModel implements InferenceModel {
  */
 export async function loadGraphModel(
     modelUrl: string|io.IOHandler,
-    options: io.LoadOptions = {}): Promise<GraphModel> {
+    options: io.LoadOptions = {}, tfio = io): Promise<GraphModel> {
   if (modelUrl == null) {
     throw new Error(
         'modelUrl in loadGraphModel() cannot be null. Please provide a url ' +
@@ -456,7 +459,7 @@ export async function loadGraphModel(
       modelUrl = `${modelUrl}${DEFAULT_MODEL_NAME}${TFHUB_SEARCH_PARAM}`;
     }
   }
-  const model = new GraphModel(modelUrl, options);
+  const model = new GraphModel(modelUrl, options, tfio);
   await model.load();
   return model;
 }

tfjs-converter/src/executor/graph_model_test.ts

@@ -23,6 +23,7 @@ import {deregisterOp, registerOp} from '../operations/custom_op/register';
 import {GraphNode} from '../operations/types';
 
 import {GraphModel, loadGraphModel} from './graph_model';
+import {RecursiveSpy, spyOnAllFunctions} from '../operations/executors/spy_ops';
 
 const HOST = 'http://example.org';
 const MODEL_URL = `${HOST}/model.json`;
@@ -368,6 +369,12 @@ describe('loadSync', () => {
 });
 
 describe('loadGraphModel', () => {
+  let spyIo: RecursiveSpy<typeof io>;
+
+  beforeEach(() => {
+    spyIo = spyOnAllFunctions(io);
+  });
+
   it('Pass a custom io handler', async () => {
     const customLoader: tfc.io.IOHandler = {
       load: async () => {
@@ -397,22 +404,25 @@ describe('loadGraphModel', () => {
 
   it('Pass a fetchFunc', async () => {
     const fetchFunc = () => {};
-    spyOn(tfc.io, 'getLoadHandlers').and.returnValue([
+    spyIo.getLoadHandlers.and.returnValue([
       CUSTOM_HTTP_MODEL_LOADER
     ]);
-    await loadGraphModel(MODEL_URL, {fetchFunc});
-    expect(tfc.io.getLoadHandlers).toHaveBeenCalledWith(MODEL_URL, {fetchFunc});
+    await loadGraphModel(MODEL_URL, {fetchFunc}, spyIo);
+    expect(spyIo.getLoadHandlers).toHaveBeenCalledWith(MODEL_URL, {fetchFunc});
   });
 });
 
 describe('Model', () => {
+  let spyIo: RecursiveSpy<typeof io>;
+
   beforeEach(() => {
-    model = new GraphModel(MODEL_URL);
+    spyIo = spyOnAllFunctions(io);
+    model = new GraphModel(MODEL_URL, undefined, spyIo);
   });
 
   describe('custom model', () => {
     beforeEach(() => {
-      spyOn(tfc.io, 'getLoadHandlers').and.returnValue([
+      spyIo.getLoadHandlers.and.returnValue([
         CUSTOM_HTTP_MODEL_LOADER
       ]);
       registerOp('CustomOp', (nodeValue: GraphNode) => {
@@ -454,11 +464,10 @@ describe('Model', () => {
 
   describe('simple model', () => {
     beforeEach(() => {
-      spyOn(tfc.io, 'getLoadHandlers').and.returnValue([
+      spyIo.getLoadHandlers.and.returnValue([
         SIMPLE_HTTP_MODEL_LOADER
       ]);
-      spyOn(tfc.io, 'browserHTTPRequest')
-          .and.returnValue(SIMPLE_HTTP_MODEL_LOADER);
+      spyIo.browserHTTPRequest.and.returnValue(SIMPLE_HTTP_MODEL_LOADER);
     });
     it('load', async () => {
       const loaded = await model.load();
@@ -591,7 +600,7 @@ describe('Model', () => {
     describe('dispose', () => {
       it('should dispose the weights', async () => {
         const numOfTensors = tfc.memory().numTensors;
-        model = new GraphModel(MODEL_URL);
+        model = new GraphModel(MODEL_URL, undefined, spyIo);
 
         await model.load();
         model.dispose();
@@ -609,7 +618,7 @@ describe('Model', () => {
 
     describe('relative path', () => {
       beforeEach(() => {
-        model = new GraphModel(RELATIVE_MODEL_URL);
+        model = new GraphModel(RELATIVE_MODEL_URL, undefined, spyIo);
       });
 
       it('load', async () => {
@@ -619,22 +628,22 @@ describe('Model', () => {
     });
 
     it('should loadGraphModel', async () => {
-      const model = await loadGraphModel(MODEL_URL);
+      const model = await loadGraphModel(MODEL_URL, undefined, spyIo);
       expect(model).not.toBeUndefined();
     });
 
     it('should loadGraphModel with request options', async () => {
       const model = await loadGraphModel(
-          MODEL_URL, {requestInit: {credentials: 'include'}});
-      expect(tfc.io.browserHTTPRequest).toHaveBeenCalledWith(MODEL_URL, {
+          MODEL_URL, {requestInit: {credentials: 'include'}}, spyIo);
+      expect(spyIo.browserHTTPRequest).toHaveBeenCalledWith(MODEL_URL, {
         requestInit: {credentials: 'include'}
       });
       expect(model).not.toBeUndefined();
     });
 
     it('should call loadGraphModel for TfHub Module', async () => {
       const url = `${HOST}/model/1`;
-      const model = await loadGraphModel(url, {fromTFHub: true});
+      const model = await loadGraphModel(url, {fromTFHub: true}, spyIo);
       expect(model).toBeDefined();
     });
 
@@ -656,11 +665,10 @@ describe('Model', () => {
 
   describe('control flow model', () => {
     beforeEach(() => {
-      spyOn(tfc.io, 'getLoadHandlers').and.returnValue([
+      spyIo.getLoadHandlers.and.returnValue([
         CONTROL_FLOW_HTTP_MODEL_LOADER
       ]);
-      spyOn(tfc.io, 'browserHTTPRequest')
-          .and.returnValue(CONTROL_FLOW_HTTP_MODEL_LOADER);
+      spyIo.browserHTTPRequest.and.returnValue(CONTROL_FLOW_HTTP_MODEL_LOADER);
     });
 
     describe('save', () => {
@@ -747,11 +755,10 @@ describe('Model', () => {
   };
   describe('dynamic shape model', () => {
     beforeEach(() => {
-      spyOn(tfc.io, 'getLoadHandlers').and.returnValue([
+      spyIo.getLoadHandlers.and.returnValue([
         DYNAMIC_HTTP_MODEL_LOADER
       ]);
-      spyOn(tfc.io, 'browserHTTPRequest')
-          .and.returnValue(DYNAMIC_HTTP_MODEL_LOADER);
+      spyIo.browserHTTPRequest.and.returnValue(DYNAMIC_HTTP_MODEL_LOADER);
     });
 
     it('should throw error if call predict directly', async () => {
@@ -792,11 +799,10 @@ describe('Model', () => {
   });
   describe('dynamic shape model with metadata', () => {
     beforeEach(() => {
-      spyOn(tfc.io, 'getLoadHandlers').and.returnValue([
+      spyIo.getLoadHandlers.and.returnValue([
         DYNAMIC_HTTP_MODEL_NEW_LOADER
       ]);
-      spyOn(tfc.io, 'browserHTTPRequest')
-          .and.returnValue(DYNAMIC_HTTP_MODEL_NEW_LOADER);
+      spyIo.browserHTTPRequest.and.returnValue(DYNAMIC_HTTP_MODEL_NEW_LOADER);
     });
 
     it('should be success if call executeAsync with signature key',
@@ -818,11 +824,10 @@ describe('Model', () => {
 
   describe('Hashtable model', () => {
     beforeEach(() => {
-      spyOn(tfc.io, 'getLoadHandlers').and.returnValue([
+      spyIo.getLoadHandlers.and.returnValue([
         HASHTABLE_HTTP_MODEL_LOADER
       ]);
-      spyOn(tfc.io, 'browserHTTPRequest')
-          .and.returnValue(HASHTABLE_HTTP_MODEL_LOADER);
+      spyIo.browserHTTPRequest.and.returnValue(HASHTABLE_HTTP_MODEL_LOADER);
     });
     it('should be successful if call executeAsync', async () => {
       await model.load();