@@ -379,6 +379,8 @@ class TransformWithStateInitialStateSuite extends StateStoreMetricsTest
379379 withSQLConf(SQLConf .STATE_STORE_PROVIDER_CLASS .key ->
380380 classOf [RocksDBStateStoreProvider ].getName) {
381381
382+ val clock = new StreamManualClock
383+
382384 val inputData = MemoryStream [InitInputRow ]
383385 val kvDataSet = inputData.toDS()
384386 .groupByKey(x => x.key)
@@ -390,10 +392,12 @@ class TransformWithStateInitialStateSuite extends StateStoreMetricsTest
390392 TimeMode .None (), OutputMode .Append (), initStateDf)
391393
392394 testStream(query, OutputMode .Update ())(
395+ StartStream (Trigger .ProcessingTime (" 1 second" ), triggerClock = clock),
393396 // non-exist key test
394397 AddData (inputData, InitInputRow (" k1" , " update" , 37.0 )),
395398 AddData (inputData, InitInputRow (" k2" , " update" , 40.0 )),
396399 AddData (inputData, InitInputRow (" non-exist" , " getOption" , - 1.0 )),
400+ AdvanceManualClock (1 * 1000 ),
397401 CheckNewAnswer ((" non-exist" , " getOption" , - 1.0 )),
398402 Execute { q =>
399403 assert(q.lastProgress
@@ -402,59 +406,80 @@ class TransformWithStateInitialStateSuite extends StateStoreMetricsTest
402406 AddData (inputData, InitInputRow (" k1" , " appendList" , 37.0 )),
403407 AddData (inputData, InitInputRow (" k2" , " appendList" , 40.0 )),
404408 AddData (inputData, InitInputRow (" non-exist" , " getList" , - 1.0 )),
409+ AdvanceManualClock (1 * 1000 ),
405410 CheckNewAnswer (),
406411
407412 AddData (inputData, InitInputRow (" k1" , " incCount" , 37.0 )),
408413 AddData (inputData, InitInputRow (" k2" , " incCount" , 40.0 )),
409414 AddData (inputData, InitInputRow (" non-exist" , " getCount" , - 1.0 )),
415+ AdvanceManualClock (1 * 1000 ),
410416 CheckNewAnswer ((" non-exist" , " getCount" , 0.0 )),
417+
411418 AddData (inputData, InitInputRow (" k2" , " incCount" , 40.0 )),
412419 AddData (inputData, InitInputRow (" k2" , " getCount" , 40.0 )),
420+ AdvanceManualClock (1 * 1000 ),
413421 CheckNewAnswer ((" k2" , " getCount" , 2.0 )),
414422
415423 // test every row in initial State is processed
416424 AddData (inputData, InitInputRow (" init_1" , " getOption" , - 1.0 )),
425+ AdvanceManualClock (1 * 1000 ),
417426 CheckNewAnswer ((" init_1" , " getOption" , 40.0 )),
427+
418428 AddData (inputData, InitInputRow (" init_2" , " getOption" , - 1.0 )),
429+ AdvanceManualClock (1 * 1000 ),
419430 CheckNewAnswer ((" init_2" , " getOption" , 100.0 )),
420431
421432 AddData (inputData, InitInputRow (" init_1" , " getList" , - 1.0 )),
433+ AdvanceManualClock (1 * 1000 ),
422434 CheckNewAnswer ((" init_1" , " getList" , 40.0 )),
435+
423436 AddData (inputData, InitInputRow (" init_2" , " getList" , - 1.0 )),
437+ AdvanceManualClock (1 * 1000 ),
424438 CheckNewAnswer ((" init_2" , " getList" , 100.0 )),
425439
426440 AddData (inputData, InitInputRow (" init_1" , " getCount" , 40.0 )),
441+ AdvanceManualClock (1 * 1000 ),
427442 CheckNewAnswer ((" init_1" , " getCount" , 1.0 )),
443+
428444 AddData (inputData, InitInputRow (" init_2" , " getCount" , 100.0 )),
445+ AdvanceManualClock (1 * 1000 ),
429446 CheckNewAnswer ((" init_2" , " getCount" , 1.0 )),
430447
431448 // Update row with key in initial row will work
432449 AddData (inputData, InitInputRow (" init_1" , " update" , 50.0 )),
433450 AddData (inputData, InitInputRow (" init_1" , " getOption" , - 1.0 )),
451+ AdvanceManualClock (1 * 1000 ),
434452 CheckNewAnswer ((" init_1" , " getOption" , 50.0 )),
453+
435454 AddData (inputData, InitInputRow (" init_1" , " remove" , - 1.0 )),
436455 AddData (inputData, InitInputRow (" init_1" , " getOption" , - 1.0 )),
456+ AdvanceManualClock (1 * 1000 ),
437457 CheckNewAnswer ((" init_1" , " getOption" , - 1.0 )),
438458
439459 AddData (inputData, InitInputRow (" init_1" , " appendList" , 50.0 )),
440460 AddData (inputData, InitInputRow (" init_1" , " getList" , - 1.0 )),
461+ AdvanceManualClock (1 * 1000 ),
441462 CheckNewAnswer ((" init_1" , " getList" , 50.0 ), (" init_1" , " getList" , 40.0 )),
442463
443464 AddData (inputData, InitInputRow (" init_1" , " incCount" , 40.0 )),
444465 AddData (inputData, InitInputRow (" init_1" , " getCount" , 40.0 )),
466+ AdvanceManualClock (1 * 1000 ),
445467 CheckNewAnswer ((" init_1" , " getCount" , 2.0 )),
446468
447469 // test remove
448470 AddData (inputData, InitInputRow (" k1" , " remove" , - 1.0 )),
449471 AddData (inputData, InitInputRow (" k1" , " getOption" , - 1.0 )),
472+ AdvanceManualClock (1 * 1000 ),
450473 CheckNewAnswer ((" k1" , " getOption" , - 1.0 )),
451474
452475 AddData (inputData, InitInputRow (" init_1" , " clearCount" , - 1.0 )),
453476 AddData (inputData, InitInputRow (" init_1" , " getCount" , - 1.0 )),
477+ AdvanceManualClock (1 * 1000 ),
454478 CheckNewAnswer ((" init_1" , " getCount" , 0.0 )),
455479
456480 AddData (inputData, InitInputRow (" init_1" , " clearList" , - 1.0 )),
457481 AddData (inputData, InitInputRow (" init_1" , " getList" , - 1.0 )),
482+ AdvanceManualClock (1 * 1000 ),
458483 CheckNewAnswer ()
459484 )
460485 }
0 commit comments