Skip to content

Commit fff6793

Browse files
anishshri-dbHeartSaVioR
authored andcommitted
[SPARK-50516][SS][MINOR] Fix the init state related test to use StreamManualClock
### What changes were proposed in this pull request? Fix the init state related test to use StreamManualClock ### Why are the changes needed? Fix flakiness on CI runs ``` [info] Run completed in 2 minutes, 43 seconds. [info] Total number of tests run: 22 [info] Suites: completed 1, aborted 0 [info] Tests: succeeded 22, failed 0, canceled 0, ignored 0, pending 0 [info] All tests passed. ``` ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Test only change ### Was this patch authored or co-authored using generative AI tooling? No Closes #49102 from anishshri-db/task/SPARK-50516. Authored-by: Anish Shrigondekar <[email protected]> Signed-off-by: Jungtaek Lim <[email protected]>
1 parent 2fea84e commit fff6793

File tree

1 file changed

+25
-0
lines changed

1 file changed

+25
-0
lines changed

sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateInitialStateSuite.scala

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)