Merge pull request #340 from code-dot-org/more-tests

More tests

Author: bethanyaconnor

src/oceans/models/predict.js

@@ -29,28 +29,30 @@ export const init = () => {
   // seconds, and if it happens immediately it will prevent React from rendering
   // the loading UI first.  If we are going to show the loading spinner, then also
   // delay the beginning of our training.
-  setTimeout(() => {
-    // Manually send a GA event for Modes.Predicting.
-    reportPageView('predicting');
+  return new Promise(resolve => setTimeout(resolve, trainingDelayTime)).then(
+    () => {
+      // Manually send a GA event for Modes.Predicting.
+      reportPageView('predicting');
 
-    state.trainer.train();
+      state.trainer.train();
 
-    let fishData = [];
-    if (state.appMode === AppMode.CreaturesVTrashDemo) {
-      fishData = fishData.concat(generateOcean(4, 0, true, true, false));
-      fishData = fishData.concat(generateOcean(3, 4, false, false, true));
-    } else if (state.appMode === AppMode.FishLong) {
-      fishData = generateOcean(500);
-    } else {
-      fishData = generateOcean(100);
-    }
+      let fishData = [];
+      if (state.appMode === AppMode.CreaturesVTrashDemo) {
+        fishData = fishData.concat(generateOcean(4, 0, true, true, false));
+        fishData = fishData.concat(generateOcean(3, 4, false, false, true));
+      } else if (state.appMode === AppMode.FishLong) {
+        fishData = generateOcean(500);
+      } else {
+        fishData = generateOcean(100);
+      }
 
-    if (setLoadingSpinner) {
-      finishLoading(startTime, () => onLoadComplete(fishData));
-    } else {
-      onLoadComplete(fishData);
+      if (setLoadingSpinner) {
+        finishLoading(startTime, () => onLoadComplete(fishData));
+      } else {
+        onLoadComplete(fishData);
+      }
     }
-  }, trainingDelayTime);
+  );
 };
 
 const onLoadComplete = fishData => {
@@ -63,8 +65,7 @@ const onLoadComplete = fishData => {
 export const predictFish = (state, idx) => {
   return new Promise(resolve => {
     const fish = state.fishData[idx];
-
-    state.trainer.predict(fish).then(prediction => {
+    return state.trainer.predict(fish).then(prediction => {
       fish.setResult(prediction);
       resolve(prediction);
     });

test/unit/oceans/models/predict.test.js

@@ -1,45 +1,61 @@
 const {initFishData} = require('@ml/utils/fishData');
 import {setState, getState, resetState} from '@ml/oceans/state';
-import {ClassType, Modes} from '@ml/oceans/constants';
+import {ClassType, Modes, AppMode} from '@ml/oceans/constants';
 import {init, predictFish} from '@ml/oceans/models/predict';
 import SimpleTrainer from '@ml/utils/SimpleTrainer';
+import {TrashOceanObject} from '@ml/oceans/OceanObject';
 
-describe('Model quality test', () => {
+describe('Predict test', () => {
   beforeAll(() => {
     initFishData();
+    jasmine.DEFAULT_TIMEOUT_INTERVAL = 5000;
   });
 
   beforeEach(() => {
+    const trainer = new SimpleTrainer();
+    trainer.predict = jest.fn(async example => {
+      if (example instanceof TrashOceanObject) {
+        return {predictedClassId: 1, confidenceByClassId: {0: 0, 1: 1}};
+      } else {
+        return {predictedClassId: 0, confidenceByClassId: {0: 1, 1: 0}};
+      }
+    });
+
     resetState();
     setState({
-      trainer: new SimpleTrainer()
+      trainer
     });
   });
 
-  jest.useFakeTimers();
-
   test('init state without intermediate loading', () => {
-    init();
-    // There's a setTimeout in the predict init. Ideally this would be
-    // refactored to use promises if possible but until then, wait 100 ms.
-    setTimeout(async () => {
+    return init().then(() => {
       const state = getState();
       expect(state).toBeTruthy();
-      expect(state.mode).toBe(Modes.Predicting);
+      expect(state.currentMode).toBe(Modes.Predicting);
       expect(state.fishData).toBeTruthy();
       expect(state.fishData.length).toBeGreaterThan(0);
-    }, 100);
+    });
   });
 
-  test('fish gets prediction on predict', async () => {
+  test('init state to intermediate loading', async () => {
+    setState({appMode: AppMode.FishShort});
     init();
-    // There's a setTimeout in the predict init. Ideally this would be
-    // refactored to use promises if possible but until then, wait 100 ms.
-    setTimeout(async () => {
-      const state = getState();
-      await predictFish(state, 0);
-      expect(state.trainer.predict).toBeCalled();
-      expect(state.fishData[0].result).toBeInstanceOf(Number);
-    }, 100);
+
+    const state = getState();
+    expect(state).toBeTruthy();
+    expect(state.currentMode).toBe(Modes.IntermediateLoading);
+  });
+
+  it('fish gets prediction on predict', async () => {
+    setState({appMode: AppMode.FishVTrash});
+    let state = getState();
+    return init().then(() => {
+      state = getState();
+      return predictFish(state, 0).then(() => {
+        expect(
+          state.fishData[0].result.predictedClassId
+        ).toBeGreaterThanOrEqual(0);
+      });
+    });
   });
 });