Flaky test - TestStatefulTask#testShouldStartAndRestore

testShouldStartAndRestore FAILED
org.junit.ComparisonFailure: expected:<[4]> but was:<[99]>
at org.junit.Assert.assertEquals(Assert.java:123)
at org.junit.Assert.assertEquals(Assert.java:145)
at org.apache.samza.test.integration.StreamTaskTestUtil.send(StreamTaskTestUtil.scala:235)
at org.apache.samza.test.integration.TestStatefulTask.testShouldRestoreStore(TestStatefulTask.scala:134)
at org.apache.samza.test.integration.TestStatefulTask.testShouldStartAndRestore(TestStatefulTask.scala:77)

https://builds.apache.org/job/samza-pr-checks/824/consoleFull

I think it's because of this commit, which increased the number of messages produced to the topic but didn't adjust the expected number of messages to restore and the code has a race condition if doesn't wait for all the messages to be consumed.
https://github.com/apache/samza/commit/12594fb710260b62f9a21a8be785bd8dd5dcdd01