Make KinesisSourceGraphIntegrationSpec more robust. (#46)

* Make KinesisSourceGraphIntegrationSpec more robust.

All test cases assumed exactly once delivery - now assume at least once, which is what Kinesis provides.
Let the spec operate on one shard, so we are able to specify conditions on the content easily and use `takeWhile` instead of `take`.
Specify the batchSize exactly, where the default does not apply.
Remove the 4 worker test, since more worker than shards is already covered in the 2 worker setup.

Author: aquamatthias

src/it/scala/com/weightwatchers/reactive/kinesis/common/KinesisSuite.scala

@@ -152,7 +152,7 @@ trait KinesisSuite
 
   def TestStreamNrOfMessagesPerShard: Long
 
-  def TestStreamNumberOfShards: Long = 2
+  def TestStreamNumberOfShards: Long = 1
 
   private lazy val kclSetupConfig =
     consumerConfFor(streamName = TestStreamName, appName = suiteName).kclConfiguration

src/it/scala/com/weightwatchers/reactive/kinesis/consumer/ConsumerProcessingManagerIntegrationSpec.scala

@@ -27,6 +27,7 @@ class ConsumerProcessingManagerIntegrationSpec
     with Eventually {
 
   override def TestStreamNrOfMessagesPerShard: Long    = 0
+  override def TestStreamNumberOfShards: Long          = 2
   override implicit def patienceConfig: PatienceConfig = PatienceConfig(60.seconds, 1.second)
 
   "A ConsumerProcessingManager on a stream with 2 shards" - {

src/it/scala/com/weightwatchers/reactive/kinesis/stream/KinesisSourceGraphIntegrationSpec.scala

@@ -26,17 +26,18 @@ class KinesisSourceGraphIntegrationSpec
     "process all messages of a stream with one worker" in new withKinesisConfForApp("1worker") {
       val result = Kinesis
         .source(consumerConf = consumerConf())
-        .take(TestStreamNumberOfShards * TestStreamNrOfMessagesPerShard)
+        .takeWhile(_.payload.payloadAsString().toLong < TestStreamNrOfMessagesPerShard,
+                   inclusive = true)
         .map { event =>
           event.commit()
-          event.payload.payload
+          event.payload.payloadAsString()
         }
         .runWith(Sink.seq)
 
       val grouped = result.futureValue.groupBy(identity)
       result.futureValue.distinct should have size TestStreamNrOfMessagesPerShard
       grouped should have size TestStreamNrOfMessagesPerShard
-      grouped.values.foreach(_ should have size TestStreamNumberOfShards)
+      grouped.values.foreach(_.size.toLong shouldBe >=(TestStreamNumberOfShards))
     }
 
     "process all messages of a stream with 2 workers" in new withKinesisConfForApp("2worker") {
@@ -48,44 +49,18 @@ class KinesisSourceGraphIntegrationSpec
       val source2   = Kinesis.source(consumerConf = consumerConf())
       val result = source1
         .merge(source2)
-        .take(TestStreamNrOfMessagesPerShard * TestStreamNumberOfShards)
+        .takeWhile(_.payload.payloadAsString().toLong < TestStreamNrOfMessagesPerShard,
+                   inclusive = true)
         .map { event =>
           event.commit()
-          event.payload.payload
+          event.payload.payloadAsString()
         }
         .runWith(Sink.seq)
 
       val grouped = result.futureValue.groupBy(identity)
       result.futureValue.distinct should have size TestStreamNrOfMessagesPerShard
       grouped should have size TestStreamNrOfMessagesPerShard
-      grouped.values.foreach(_ should have size TestStreamNumberOfShards)
-    }
-
-    "process all messages of a stream with 4 workers" in new withKinesisConfForApp("4worker") {
-      // Please note: since all sources are started simultaneously, all will assume there is no other worker.
-      // During register all except one will fail and not read any message until retry
-      // Depending on timing one or multiple sources will read all events
-      val batchSize = TestStreamNrOfMessagesPerShard
-      val source1   = Kinesis.source(consumerConf = consumerConf())
-      val source2   = Kinesis.source(consumerConf = consumerConf())
-      val source3   = Kinesis.source(consumerConf = consumerConf())
-      val source4   = Kinesis.source(consumerConf = consumerConf())
-      val result = source1
-        .merge(source2)
-        .merge(source3)
-        .merge(source4)
-        // Since only 2 clients can take batchSize messages, an overall take is needed here to end the stream
-        .take(TestStreamNrOfMessagesPerShard * TestStreamNumberOfShards)
-        .map { event =>
-          event.commit()
-          event.payload.payload
-        }
-        .runWith(Sink.seq)
-
-      val grouped = result.futureValue.groupBy(identity)
-      result.futureValue.distinct should have size TestStreamNrOfMessagesPerShard
-      grouped should have size TestStreamNrOfMessagesPerShard
-      grouped.values.foreach(_ should have size TestStreamNumberOfShards)
+      grouped.values.foreach(_.size.toLong shouldBe >=(TestStreamNumberOfShards))
     }
 
     "maintain the read position in the stream correctly" in new withKinesisConfForApp(
@@ -98,15 +73,16 @@ class KinesisSourceGraphIntegrationSpec
       // - dies after one batch
       // We expect to get all messages by n reads (which means, that the read position was stored correctly)
       val result =
-        for (_ <- 1
+        for (iteration <- 1
                .to((TestStreamNumberOfShards * TestStreamNrOfMessagesPerShard / batchSize).toInt))
           yield {
             Kinesis
-              .source(consumerConf = consumerConf())
-              .take(batchSize)
+              .source(consumerConf = consumerConf(batchSize = batchSize))
+              .takeWhile(_.payload.payloadAsString().toLong < batchSize * iteration,
+                         inclusive = true)
               .map { event =>
                 event.commit()
-                event.payload.payload
+                event.payload
               }
               .runWith(Sink.seq)
               .futureValue
@@ -117,38 +93,39 @@ class KinesisSourceGraphIntegrationSpec
       val grouped = allMessages.groupBy(identity)
       allMessages.distinct should have size TestStreamNrOfMessagesPerShard
       grouped should have size TestStreamNrOfMessagesPerShard
+      grouped.values.foreach(_.size.toLong shouldBe >=(TestStreamNumberOfShards))
     }
 
     "not commit the position, if the event is not committed" in new withKinesisConfForApp(
       "not_committed"
     ) {
-      val batchSize = TestStreamNrOfMessagesPerShard / 2 // 2 * NrOfShards batches needed
-
-      // This worker will read batchSize events and will not commit
+      // This worker will read all events and will not commit
       // We expect that the read position will not change
       val uncommitted = Kinesis
         .source(consumerConf())
-        .take(batchSize)
+        .takeWhile(_.payload.payloadAsString().toLong < TestStreamNrOfMessagesPerShard,
+                   inclusive = true)
         .runWith(Sink.seq)
         .futureValue
 
       // This worker will read all available events.
       // This works only, if the first worker has not committed anything
       val committed = Kinesis
         .source(consumerConf = consumerConf())
-        .take(TestStreamNumberOfShards * TestStreamNrOfMessagesPerShard)
+        .takeWhile(_.payload.payloadAsString().toLong < TestStreamNrOfMessagesPerShard,
+                   inclusive = true)
         .map { event =>
           event.commit()
-          event.payload.payload
+          event.payload.payloadAsString()
         }
         .runWith(Sink.seq)
         .futureValue
 
-      uncommitted should have size batchSize
+      uncommitted should have size TestStreamNrOfMessagesPerShard
       val grouped = committed.groupBy(identity)
       committed.distinct should have size TestStreamNrOfMessagesPerShard
       grouped should have size TestStreamNrOfMessagesPerShard
-      grouped.values.foreach(_ should have size TestStreamNumberOfShards)
+      grouped.values.foreach(_.size.toLong shouldBe >=(TestStreamNumberOfShards))
     }
   }
 }