Merge pull request #1031 from stephenplusplus/spp--prediction-error

prediction: allow more time for model training
googleapis · Dec 19, 2015 · f919245 · f919245
2 parents 2f05c2f + 136d50a
commit f919245
Showing 1 changed file with 31 additions and 28 deletions.
diff --git a/system-test/prediction.js b/system-test/prediction.js
@@ -193,44 +193,47 @@ describe('Prediction', function() {
       before(function(done) {
         model.create({
           data: file
-        }, done);
+        }, function(err) {
+          if (err) {
+            done(err);
+            return;
+          }
+
+          function isModelTrained(callback) {
+            model.getMetadata(function(err, metadata) {
+              if (err) {
+                callback(err);
+                return;
+              }
+
+              if (metadata.trainingStatus === 'RUNNING') {
+                callback(new Error('Model still training.'));
+                return;
+              }
+
+              callback();
+            });
+          }
+
+          async.retry({ times: 5, interval: 10000 }, isModelTrained, done);
+        });
       });
 
       after(function(done) {
         model.delete(done);
       });
 
       it('should return results', function(done) {
-        // The model needs to process the new data before it can be queried.
-        async.retry(
-          {
-            times: 5,
-            interval: 7500
-          },
-
-          function(callback) {
-            model.query('hello', callback);
-          },
-
-          function(err, results) {
-            assert.ifError(err);
-            assert.strictEqual(results.winner, 'English');
-            assert.strictEqual(results.scores[0].score, 1);
-            done();
-          });
+        model.query('hello', function(err, results) {
+          assert.ifError(err);
+          assert.strictEqual(results.winner, 'English');
+          assert.strictEqual(results.scores[0].score, 1);
+          done();
+        });
       });
 
       it('should return analysis', function(done) {
-        var model = prediction.model(generateName());
-
-        model.create(function(err, model) {
-          assert.ifError(err);
-
-          model.analyze(function(err) {
-            assert.ifError(err);
-            model.delete(done);
-          });
-        });
+        model.analyze(done);
       });
     });
   });