combinator to await the n-th element or the first n elements of a stream

[Angel-O]

Quick look TL;DR

A simple api to await the n-th element of a stream. (can also await the first n elements. See code at the end)
The combinator will return the n-th element or the first n soon as they are available.

    *   // infinite stream 10, 20, 30, 40...
    *   val streamA = Stream.unfold(10)(s => Some((s, s + 10))).covary[IO]
    *
    *   // finite stream
    *   val streamB = Stream(10,20,30,40,50).covary[IO]
    *
    *   val thirdElementOfA = await(streamA).nTH[3]
    *   val first2ElementsOfB = await(streamB).firstN[3]
    *   val tenthElementOfB = await(streamB).nTH[10]
    *
    *   thirdElementOfA.unsafeRunSync() // returns 30 eventually
    *   first2ElementsOfB.unsafeRunSync() // returns List(10, 20) eventually
    *   tenthElementsOfB.unsafeRunSync() // never produces a result
    *
    *   // shorter syntax (if you are happy with the default interval)
    *
    *   import some.package.StreamUtils._
    *
    *   val stream: Stream[IO, A] = ???
    *
    *   stream.awaitNth[5]
    *   stream.awaitFirstN[5]

Context

While working on an integration test I noticed the test logic was a bit heavy so I decided to declutter it

Use case

IT test that publishes events through a kafka pipeline and stores events to a Postgres db.
A test consumer is subscribed to one of the topics so that we can inspect the behaviour of the pipeline and make assertions at the end of the test

Issues

• calls to .delayBy(xyz.seconds) to make sure the events have enough time to flow through the pipeline before ending up to the database
• inaccuracy when setting expectations: Stream.take(n) does not enforce the correct number of elements
• boilerplate code to create Deferred

At a very high level the old test looks like this. We publish events to a topic, we read from 2 streams (finite or infinite) and we complete 2 deferred to make assertions on certain events. There are more components in the original test, removed for simplicity

val dbStream = db.getEvents
val consumerRecords = consumer.getRecords
val awaitDbEvent: Deferred[F, A] =  ...
val awaitConsumerRecord: Deferred[F, B] = ...
val n = 3 // number of events expected on the database
val m = 2 // number of events expected on the topic

for {
_        <- produceEvents
_        <- dbStream.take(n).compile.toList.flatMap(as => awaitDbEvent.complete(as.last)).delayBy(2.seconds) // wait for events to reach db
_        <- consumerRecords.take(m).compile.toList.flatMap(as => awaitConsumerRecord.complete(as.last))
a        <- awaitDbEvent.get
b        <- awaitConsumerRecord.get
} yield {
assert(a, ...)
assert(b, ...)
}

The new version using the awaitNth[N] combinator

val dbStream = db.getEvents
val consumerRecords = consumer.getRecords

for {
_        <- produceEvents
a        <- dbStream.awaitNth[3] // no need to delay 
b        <- consumerRecords.awaitNth[2] // unlike Stream.take(n) if we set a higher expectation the test will fail (we need to be precise)
} yield {
assert(a, ...)
assert(b, ...)
}

Advantages

• no need to declare expectations on number of event (if you don't like magic numbers in your code, you usually store these as values)
• decluttering test code from boilerplate and low level details (Deferred) to focus on the business logic being tested
• Stream.take(n) is quite permissive in the sense that it will accept a number greater than the number of events that the stream can produce. If you set the wrong expectation, the test might still pass: not ideal in the context of a test. The awaitNth combinator will wait for the exact number of elements to be produced.
• no need to add delays for finite streams (in the example above we would need to add a delay to make sure that the db stream doesn't run before the records flowing through the kafka pipeline reach the database. Could be solved by retrieving the kafka records before the db events, but it makes the test more brittle)
• related to the point above: higher degree of flexibility when it comes to refactoring (it doesn't matter if the stream is finite or not, the order does not matter: at least in my use case :)

Notes

• probably the original test could have been written in a simpler way
• the combinator is agnostic wrt the stream being finite or infinite (is it an advantage or a disadvantage 🤔 )

Critiques

• it is quite a specific use case (e.g. makes assumptions about the finite streams not having side effects)
• the optional interval param might not be great

The code

package some.package

import cats.effect.{Deferred, Temporal}
import cats.implicits.*
import fs2.Stream

import scala.concurrent.duration.{DurationInt, FiniteDuration}

object StreamUtils {

  class StreamAwaitPartiallyApplied[F[_]: Temporal, A] private[StreamUtils] (stream: Stream[F, A], interval: FiniteDuration) {

    def nTH[N <: Int: ValueOf]: F[A] =
      runStream[Int, A](0, (acc, cur) => (acc + 1, Option.when(acc == valueOf[N] - 1)(cur)))

    def firstN[N <: Int: ValueOf]: F[List[A]] =
      runStream[(Int, List[A]), List[A]](
        (0, Nil),
        { case ((count, as), a) => ((count + 1, a +: as), Option.when(count == valueOf[N] - 1)(a +: as)) }
      ).map(_.reverse)

    private def runStream[Z, O](zero: Z, run: (Z, A) => (Z, Option[O])): F[O] =
      Deferred[F, O].flatMap { deferred =>
        val writeO = stream.mapAccumulate(zero)(run).evalTap { case (_, maybeO) => maybeO.traverse_(deferred.complete) }
        val readO  = deferred.get
        Stream.eval(readO).concurrently((writeO ++ Stream.sleep_(interval)).repeat).compile.lastOrError
      }
  }

  def await[F[_]: Temporal, A](
      stream: Stream[F, A],
      interval: FiniteDuration = 1.second
  ): StreamAwaitPartiallyApplied[F, A] =
    new StreamAwaitPartiallyApplied(stream, interval)

  implicit class AwaitOps[F[_], A](val stream: Stream[F, A]) extends AnyVal {

    def awaitNth[N <: Int: ValueOf](implicit T: Temporal[F]): F[A] =
      await(stream).nTH

    def awaitFirstN[N <: Int: ValueOf](implicit T: Temporal[F]): F[List[A]] =
      await(stream).firstN
  }
}

[Angel-O]

• probably the original test could have been written in a simpler way

val dbStream = db.getEvents
val consumerRecords = consumer.getRecords
val n = 3 // number of events expected on the database
val m = 2 // number of events expected on the topic

for {
_       <- produceEvents
a       <- dbStream.take(n).compile.lastOrError.delayBy(2.seconds) // wait for events to reach db
b       <- consumerRecords.take(m).compile.lastOrError
} yield {
assert(a, ...)
assert(b, ...)
}

The original test also had some interrupt logic to stop a bunch of streams based on the completion status of the deferred, but as suggested by @SystemFw there's no need to use a Deferred, this version is already much cleaner than the original one: I guess the differences are:

• the behaviour of stream.await[N] vs stream.take(n), that is being more or less strict about the elements pulled from the stream (although this truly makes a difference only for streams, finite or infinite, that emit a fixed number of elements)
• the repeat behaviour, which only makes a difference for finite streams where the underlying source gets updated (e.g. a fetching entries from a Postgres db)
• the flexibility of reordering the 2nd and 3rd generators in the for-comprehension
• need to add a delay to the second generator (based on finite stream), assuming it comes right after the _ <- produceEvents generator or a retry mechanism, which by the way awaitNth[N] does as well under the hood.

which are very specific use cases (tests use cases in particular), so definitely not worth adding a combinator to the library :)

Nonetheless, It was a good exercise.

[Angel-O]

Turns out the runStream function can be simplified further: no need for a Deferred

private def runStream[Z, O](zero: Z, run: (Z, A) => (Z, Option[O])): F[O] =
  (stream
    .mapAccumulate(zero)(run)
    .map { case (_, maybeO) => maybeO } ++ Stream.sleep_(interval)).repeat.unNone.head.compile.lastOrError