There is an hidden assumption in this test case that the created TEST-TOPIC-A and TEST-TOPIC-B are propagated to the streams client at the same time, and stored as assignedTopicPartitions. However this is not always true since these two topics may be added on the client side as two consecutive metadata refreshes.
The proposed fix includes the following:
1. In waitForCondition do not trigger the conditionMet function again after the while loop, but just remember the returned value from the last call. This is safer so that if the condition changes after the while loop it will not be considered as well.
2. Not remembering a map of all the previous assigned partitions, but only the most recent one. And also get rid of the final check after streams client is closed by just use equals in the condition to make sure that it is exactly the same to the expected assignment.