@@ -381,8 +381,8 @@ def test_project_upload_dataset(self):
381381 ],
382382 "params" : {"is_prediction" : True },
383383 "assertions" : {
384- "upload" : {"count" : 2 , "kwargs" : { "is_prediction" : True } },
385- "save_annotation" : {"count" : 2 },
384+ "upload" : {"count" : 2 },
385+ "save_annotation" : {"count" : 2 , "kwargs" : { "is_prediction" : True } },
386386 },
387387 },
388388 {
@@ -392,8 +392,8 @@ def test_project_upload_dataset(self):
392392 ],
393393 "params" : {"is_prediction" : False },
394394 "assertions" : {
395- "upload" : {"count" : 1 , "kwargs" : { "is_prediction" : False } },
396- "save_annotation" : {"count" : 1 },
395+ "upload" : {"count" : 1 },
396+ "save_annotation" : {"count" : 1 , "kwargs" : { "is_prediction" : False } },
397397 },
398398 },
399399 {
@@ -410,11 +410,18 @@ def test_project_upload_dataset(self):
410410 "upload" : {
411411 "count" : 1 ,
412412 "kwargs" : {
413- "is_prediction" : True ,
414413 "batch_name" : "prediction-batch" ,
415414 "num_retry_uploads" : 2 ,
416415 },
417416 },
417+ "save_annotation" : {
418+ "count" : 1 ,
419+ "kwargs" : {
420+ "is_prediction" : True ,
421+ "job_name" : "prediction-batch" ,
422+ "num_retry_uploads" : 2 ,
423+ },
424+ },
418425 },
419426 },
420427 ]
0 commit comments