[SPARK-36109][SS][TEST] Check data after adding data to topic in KafkaSourceStressSuite

### What changes were proposed in this pull request?

This patch proposes to check data after adding data to topic in `KafkaSourceStressSuite`.

### Why are the changes needed?

The test logic in `KafkaSourceStressSuite` is not stable. For example, https://github.com/apache/spark/runs/3049244904.

Once we add data to a topic and then delete the topic before checking data, the expected answer is different to retrieved data from the sink.

### Does this PR introduce _any_ user-facing change?

No

### How was this patch tested?

Existing tests.

Closes #33311 from viirya/stream-assert.

Authored-by: Liang-Chi Hsieh <viirya@gmail.com>
Signed-off-by: Dongjoon Hyun <dongjoon@apache.org>
This commit is contained in:
Liang-Chi Hsieh 2021-07-13 01:21:32 -07:00 committed by Dongjoon Hyun
parent c46342e3d0
commit 201566cdd5
2 changed files with 8 additions and 3 deletions

View file

@ -2429,8 +2429,9 @@ class KafkaSourceStressSuite extends KafkaSourceTest {
(d, running) => { (d, running) => {
Random.nextInt(5) match { Random.nextInt(5) match {
case 0 => // Add a new topic case 0 => // Add a new topic
topics = topics ++ Seq(newStressTopic) val newTopic = newStressTopic
AddKafkaData(topics.toSet, d: _*)(message = s"Add topic $newStressTopic", topics = topics ++ Seq(newTopic)
AddKafkaData(topics.toSet, d: _*)(message = s"Add topic $newTopic",
topicAction = (topic, partition) => { topicAction = (topic, partition) => {
if (partition.isEmpty) { if (partition.isEmpty) {
testUtils.createTopic(topic, partitions = nextInt(1, 6)) testUtils.createTopic(topic, partitions = nextInt(1, 6))

View file

@ -283,7 +283,7 @@ trait StreamTest extends QueryTest with SharedSparkSession with TimeLimits with
/** Assert that a condition on the active query is true */ /** Assert that a condition on the active query is true */
class AssertOnQuery(val condition: StreamExecution => Boolean, val message: String) class AssertOnQuery(val condition: StreamExecution => Boolean, val message: String)
extends StreamAction { extends StreamAction with StreamMustBeRunning {
override def toString: String = s"AssertOnQuery(<condition>, $message)" override def toString: String = s"AssertOnQuery(<condition>, $message)"
} }
@ -871,6 +871,10 @@ trait StreamTest extends QueryTest with SharedSparkSession with TimeLimits with
case r if r < 0.7 => // AddData case r if r < 0.7 => // AddData
addRandomData() addRandomData()
// In some suites, e.g. `KafkaSourceStressSuite`, we delete Kafka topic in the
// `addData` closure. In the case, the topic with added data might be deleted
// before next check. So we must check data after adding data here.
addCheck()
case _ => // StopStream case _ => // StopStream
addCheck() addCheck()