Repository: numaproj/numaflow (go)

Issue:
flaky Test_Single
```
{"level":"info","ts":1674119528.026688,"logger":"numaflow","caller":"nats/nats.go:149","msg":"Connecting to nats service..."}
{"level":"info","ts":1674119528.0271647,"logger":"numaflow","caller":"publish/publisher.go:174","msg":"Publishing Heartbeat ticker started","entityID":"-0","otStore":"noop","hbStore":"noop"}
    nats_test.go:92: 
        	Error Trace:	nats_test.go:92
        	Error:      	Not equal: 
        	            	expected: 3
        	            	actual  : 0
        	Test:       	Test_Single
{"level":"info","ts":1674119529.0331376,"logger":"numaflow","caller":"nats/nats.go:241","msg":"Shutting down nats source server..."}
{"level":"info","ts":1674119529.0332167,"logger":"numaflow","caller":"publish/publisher.go:190","msg":"Closing watermark publisher","entityID":"-0","otStore":"noop","hbStore":"noop"}
{"level":"info","ts":1674119529.0333815,"logger":"numaflow","caller":"nats/nats.go:104","msg":"Nats disconnected"}
{"level":"info","ts":1674119529.0334904,"logger":"numaflow","caller":"nats/nats.go:250","msg":"Nats source server shutdown"}
--- FAIL: Test_Single (1.04s)
```

Pull request description:
chore: adding partition information to writer
PR changes include
- adding partition information to the writer that way we can get the partition id to which the writer is writing to
- introduced GetPartitionIdx() method to the BufferWriterInformation interface, so that all the writers will have partition information 
- readers already have the partition information.

also fixes #493 

Patch:
diff --git a/pkg/forward/forward.go b/pkg/forward/forward.go
index d455e123..8825a2c7 100644
--- a/pkg/forward/forward.go
+++ b/pkg/forward/forward.go
@@ -238,9 +238,9 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) {
 
 		// if the validation passed, we will publish the watermark to all the toBuffer partitions.
 		for toVertexName, toVertexBuffer := range isdf.toBuffers {
-			for index, partition := range toVertexBuffer {
+			for _, partition := range toVertexBuffer {
 				if p, ok := isdf.wmPublishers[toVertexName]; ok {
-					idlehandler.PublishIdleWatermark(ctx, partition, p, isdf.idleManager, int32(index), isdf.opts.logger, isdf.opts.vertexType, wmb.Watermark(time.UnixMilli(processorWMB.Watermark)))
+					idlehandler.PublishIdleWatermark(ctx, partition, p, isdf.idleManager, isdf.opts.logger, isdf.opts.vertexType, wmb.Watermark(time.UnixMilli(processorWMB.Watermark)))
 				}
 			}
 		}
@@ -431,7 +431,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) {
 					// use the watermark of the current read batch for the idle watermark
 					// same as read len==0 because there's no event published to the buffer
 					if p, ok := isdf.wmPublishers[bufferName]; ok {
-						idlehandler.PublishIdleWatermark(ctx, isdf.toBuffers[bufferName][index], p, isdf.idleManager, int32(index), isdf.opts.logger, isdf.opts.vertexType, processorWM)
+						idlehandler.PublishIdleWatermark(ctx, isdf.toBuffers[bufferName][index], p, isdf.idleManager, isdf.opts.logger, isdf.opts.vertexType, processorWM)
 					}
 				}
 			}
diff --git a/pkg/isb/interfaces.go b/pkg/isb/interfaces.go
index af8ec289..9fdcef46 100644
--- a/pkg/isb/interfaces.go
+++ b/pkg/isb/interfaces.go
@@ -62,13 +62,18 @@ type BufferReader interface {
 
 // BufferReaderInformation has information regarding the buffer we are reading from.
 type BufferReaderInformation interface {
+	// GetName returns the name.
 	GetName() string
+	// GetPartitionIdx returns the partition ID.
 	GetPartitionIdx() int32
 }
 
 // BufferWriterInformation has information regarding the buffer we are writing to.
 type BufferWriterInformation interface {
+	// GetName returns the name.
 	GetName() string
+	// GetPartitionIdx returns the partition ID.
+	GetPartitionIdx() int32
 }
 
 // SourceWatermarkPublisher publishes source watermarks based on a list of isb.ReadMessage
diff --git a/pkg/isb/stores/jetstream/writer.go b/pkg/isb/stores/jetstream/writer.go
index 373c79f5..389bb7fe 100644
--- a/pkg/isb/stores/jetstream/writer.go
+++ b/pkg/isb/stores/jetstream/writer.go
@@ -34,18 +34,19 @@ import (
 )
 
 type jetStreamWriter struct {
-	name    string
-	stream  string
-	subject string
-	conn    *jsclient.NatsConn
-	js      *jsclient.JetStreamContext
-	opts    *writeOptions
-	isFull  *atomic.Bool
-	log     *zap.SugaredLogger
+	name         string
+	partitionIdx int32
+	stream       string
+	subject      string
+	conn         *jsclient.NatsConn
+	js           *jsclient.JetStreamContext
+	opts         *writeOptions
+	isFull       *atomic.Bool
+	log          *zap.SugaredLogger
 }
 
 // NewJetStreamBufferWriter is used to provide a new instance of JetStreamBufferWriter
-func NewJetStreamBufferWriter(ctx context.Context, client jsclient.JetStreamClient, name, stream, subject string, opts ...WriteOption) (isb.BufferWriter, error) {
+func NewJetStreamBufferWriter(ctx context.Context, client jsclient.JetStreamClient, name, stream, subject string, partitionIdx int32, opts ...WriteOption) (isb.BufferWriter, error) {
 	o := defaultWriteOptions()
 	for _, opt := range opts {
 		if opt != nil {
@@ -66,14 +67,15 @@ func NewJetStreamBufferWriter(ctx context.Context, client jsclient.JetStreamClie
 	}
 
 	result := &jetStreamWriter{
-		name:    name,
-		stream:  stream,
-		subject: subject,
-		conn:    conn,
-		js:      js,
-		opts:    o,
-		isFull:  atomic.NewBool(true),
-		log:     logging.FromContext(ctx).With("bufferWriter", name).With("stream", stream).With("subject", subject),
+		name:         name,
+		partitionIdx: partitionIdx,
+		stream:       stream,
+		subject:      subject,
+		conn:         conn,
+		js:           js,
+		opts:         o,
+		isFull:       atomic.NewBool(true),
+		log:          logging.FromContext(ctx).With("bufferWriter", name).With("stream", stream).With("subject", subject),
 	}
 
 	go result.runStatusChecker(ctx)
@@ -144,6 +146,10 @@ func (jw *jetStreamWriter) GetName() string {
 	return jw.name
 }
 
+func (jw *jetStreamWriter) GetPartitionIdx() int32 {
+	return jw.partitionIdx
+}
+
 func (jw *jetStreamWriter) Close() error {
 	if jw.conn != nil && !jw.conn.IsClosed() {
 		jw.conn.Close()
diff --git a/pkg/isb/stores/redis/write.go b/pkg/isb/stores/redis/write.go
index 7b193782..448b90d1 100644
--- a/pkg/isb/stores/redis/write.go
+++ b/pkg/isb/stores/redis/write.go
@@ -41,9 +41,10 @@ var exactlyOnceInsertLuaScript string
 
 // BufferWrite is the write queue implementation powered by RedisClient.
 type BufferWrite struct {
-	Name   string
-	Stream string
-	Group  string
+	Name         string
+	PartitionIdx int32
+	Stream       string
+	Group        string
 	*BufferWriteInfo
 	*redisclient.RedisClient
 	redisclient.Options
@@ -65,7 +66,7 @@ type BufferWriteInfo struct {
 var _ isb.BufferWriter = (*BufferWrite)(nil)
 
 // NewBufferWrite returns a new redis queue writer.
-func NewBufferWrite(ctx context.Context, client *redisclient.RedisClient, name string, group string, opts ...redisclient.Option) isb.BufferWriter {
+func NewBufferWrite(ctx context.Context, client *redisclient.RedisClient, name string, group string, partitionIdx int32, opts ...redisclient.Option) isb.BufferWriter {
 	options := &redisclient.Options{
 		Pipelining:                true,
 		InfoRefreshInterval:       time.Second,
@@ -82,9 +83,10 @@ func NewBufferWrite(ctx context.Context, client *redisclient.RedisClient, name s
 
 	// check whether the script exists, if not then load
 	rqw := &BufferWrite{
-		Name:   name,
-		Stream: redisclient.GetRedisStreamName(name),
-		Group:  group,
+		Name:         name,
+		PartitionIdx: partitionIdx,
+		Stream:       redisclient.GetRedisStreamName(name),
+		Group:        group,
 		BufferWriteInfo: &BufferWriteInfo{
 			isFull:           atomic.NewBool(true),
 			refreshFullError: atomic.NewUint32(0),
@@ -254,6 +256,10 @@ func (bw *BufferWrite) GetName() string {
 	return bw.Name
 }
 
+func (bw *BufferWrite) GetPartitionIdx() int32 {
+	return bw.PartitionIdx
+}
+
 // GetGroupName gets the name of the consumer group.
 func (bw *BufferWrite) GetGroupName() string {
 	return bw.Group
diff --git a/pkg/isbsvc/redis_service.go b/pkg/isbsvc/redis_service.go
index e2df690c..c2b9c00b 100644
--- a/pkg/isbsvc/redis_service.go
+++ b/pkg/isbsvc/redis_service.go
@@ -126,7 +126,7 @@ func (r *isbsRedisSvc) ValidateBuffersAndBuckets(ctx context.Context, buffers, b
 // GetBufferInfo is used to provide buffer information like pending count, buffer length, has unprocessed data etc.
 func (r *isbsRedisSvc) GetBufferInfo(ctx context.Context, buffer string) (*BufferInfo, error) {
 	group := fmt.Sprintf("%s-group", buffer)
-	rqw := redis2.NewBufferWrite(ctx, redisclient.NewInClusterRedisClient(), buffer, group, redisclient.WithRefreshBufferWriteInfo(false))
+	rqw := redis2.NewBufferWrite(ctx, redisclient.NewInClusterRedisClient(), buffer, group, 0, redisclient.WithRefreshBufferWriteInfo(false))
 	var bufferWrite = rqw.(*redis2.BufferWrite)
 
 	bufferInfo := &BufferInfo{
diff --git a/pkg/reduce/data_forward.go b/pkg/reduce/data_forward.go
index d5b06b47..27f069ec 100644
--- a/pkg/reduce/data_forward.go
+++ b/pkg/reduce/data_forward.go
@@ -211,8 +211,8 @@ func (df *DataForward) forwardAChunk(ctx context.Context) {
 			// in this case, send idle watermark to all the toBuffer partitions
 			for toVertexName, toVertexBuffer := range df.toBuffers {
 				if publisher, ok := df.wmPublishers[toVertexName]; ok {
-					for index, bufferPartition := range toVertexBuffer {
-						idlehandler.PublishIdleWatermark(ctx, bufferPartition, publisher, df.idleManager, int32(index), df.log, dfv1.VertexTypeReduceUDF, wmb.Watermark(time.UnixMilli(processorWMB.Watermark)))
+					for _, bufferPartition := range toVertexBuffer {
+						idlehandler.PublishIdleWatermark(ctx, bufferPartition, publisher, df.idleManager, df.log, dfv1.VertexTypeReduceUDF, wmb.Watermark(time.UnixMilli(processorWMB.Watermark)))
 					}
 				}
 			}
@@ -231,8 +231,8 @@ func (df *DataForward) forwardAChunk(ctx context.Context) {
 				// so in this case, we publish an idle watermark
 				for toVertexName, toVertexBuffer := range df.toBuffers {
 					if publisher, ok := df.wmPublishers[toVertexName]; ok {
-						for index, bufferPartition := range toVertexBuffer {
-							idlehandler.PublishIdleWatermark(ctx, bufferPartition, publisher, df.idleManager, int32(index), df.log, dfv1.VertexTypeReduceUDF, wmb.Watermark(watermark))
+						for _, bufferPartition := range toVertexBuffer {
+							idlehandler.PublishIdleWatermark(ctx, bufferPartition, publisher, df.idleManager, df.log, dfv1.VertexTypeReduceUDF, wmb.Watermark(watermark))
 						}
 					}
 				}
@@ -362,8 +362,8 @@ func (df *DataForward) Process(ctx context.Context, messages []*isb.ReadMessage)
 			// this is to minimize watermark latency
 			for toVertexName, toVertexBuffer := range df.toBuffers {
 				if publisher, ok := df.wmPublishers[toVertexName]; ok {
-					for index, bufferPartition := range toVertexBuffer {
-						idlehandler.PublishIdleWatermark(ctx, bufferPartition, publisher, df.idleManager, int32(index), df.log, dfv1.VertexTypeReduceUDF, wmb.Watermark(watermark))
+					for _, bufferPartition := range toVertexBuffer {
+						idlehandler.PublishIdleWatermark(ctx, bufferPartition, publisher, df.idleManager, df.log, dfv1.VertexTypeReduceUDF, wmb.Watermark(watermark))
 					}
 				}
 			}
diff --git a/pkg/reduce/pnf/processandforward.go b/pkg/reduce/pnf/processandforward.go
index 4b855888..11bf26f9 100644
--- a/pkg/reduce/pnf/processandforward.go
+++ b/pkg/reduce/pnf/processandforward.go
@@ -307,7 +307,7 @@ func (p *processAndForward) publishWM(ctx context.Context, wm wmb.Watermark, wri
 		for index, activePartition := range activeWatermarkBuffers[toVertexName] {
 			if !activePartition {
 				if publisher, ok := p.wmPublishers[toVertexName]; ok {
-					idlehandler.PublishIdleWatermark(ctx, p.toBuffers[toVertexName][index], publisher, p.idleManager, int32(index), p.log, dfv1.VertexTypeReduceUDF, wm)
+					idlehandler.PublishIdleWatermark(ctx, p.toBuffers[toVertexName][index], publisher, p.idleManager, p.log, dfv1.VertexTypeReduceUDF, wm)
 				}
 			}
 		}
diff --git a/pkg/shared/idlehandler/idlehandler.go b/pkg/shared/idlehandler/idlehandler.go
index e13e7e2a..e0a0c8f5 100644
--- a/pkg/shared/idlehandler/idlehandler.go
+++ b/pkg/shared/idlehandler/idlehandler.go
@@ -29,8 +29,9 @@ import (
 
 // PublishIdleWatermark publishes a ctrl message with isb.Kind set to WMB. We only send one ctrl message when
 // we see a new WMB; later we only update the WMB without a ctrl message.
-func PublishIdleWatermark(ctx context.Context, toBufferPartition isb.BufferWriter, wmPublisher publish.Publisher, idleManager *wmb.IdleManager, toVertexPartition int32, logger *zap.SugaredLogger, vertexType dfv1.VertexType, wm wmb.Watermark) {
+func PublishIdleWatermark(ctx context.Context, toBufferPartition isb.BufferWriter, wmPublisher publish.Publisher, idleManager *wmb.IdleManager, logger *zap.SugaredLogger, vertexType dfv1.VertexType, wm wmb.Watermark) {
 	var toPartitionName = toBufferPartition.GetName()
+	var toVertexPartition = toBufferPartition.GetPartitionIdx()
 
 	if !idleManager.Exists(toPartitionName) {
 		if vertexType == dfv1.VertexTypeSink {
diff --git a/pkg/sinks/blackhole/blackhole.go b/pkg/sinks/blackhole/blackhole.go
index 08867f38..2e91a4e6 100644
--- a/pkg/sinks/blackhole/blackhole.go
+++ b/pkg/sinks/blackhole/blackhole.go
@@ -91,6 +91,12 @@ func (b *Blackhole) GetName() string {
 	return b.name
 }
 
+// GetPartitionIdx returns the partition index.
+// for sink it is always 0.
+func (b *Blackhole) GetPartitionIdx() int32 {
+	return 0
+}
+
 // IsFull returns whether sink is full, which is never true.
 func (b *Blackhole) IsFull() bool {
 	// printing can never be full
diff --git a/pkg/sinks/kafka/kafka.go b/pkg/sinks/kafka/kafka.go
index c7715503..caf18e8e 100644
--- a/pkg/sinks/kafka/kafka.go
+++ b/pkg/sinks/kafka/kafka.go
@@ -138,6 +138,12 @@ func (tk *ToKafka) GetName() string {
 	return tk.name
 }
 
+// GetPartitionIdx returns the partition index.
+// for sink it is always 0.
+func (tk *ToKafka) GetPartitionIdx() int32 {
+	return 0
+}
+
 // Write writes to the kafka topic.
 func (tk *ToKafka) Write(_ context.Context, messages []isb.Message) ([]isb.Offset, []error) {
 	errs := make([]error, len(messages))
diff --git a/pkg/sinks/logger/log.go b/pkg/sinks/logger/log.go
index 3382c733..755bc0cf 100644
--- a/pkg/sinks/logger/log.go
+++ b/pkg/sinks/logger/log.go
@@ -92,6 +92,12 @@ func (t *ToLog) GetName() string {
 	return t.name
 }
 
+// GetPartitionIdx returns the partition index.
+// for sink it is always 0.
+func (t *ToLog) GetPartitionIdx() int32 {
+	return 0
+}
+
 // IsFull returns whether logging is full, which is never true.
 func (t *ToLog) IsFull() bool {
 	// printing can never be full
diff --git a/pkg/sinks/udsink/sink.go b/pkg/sinks/udsink/sink.go
index 7d44eee1..d0bf6c22 100644
--- a/pkg/sinks/udsink/sink.go
+++ b/pkg/sinks/udsink/sink.go
@@ -92,6 +92,12 @@ func (s *UserDefinedSink) GetName() string {
 	return s.name
 }
 
+// GetPartitionIdx returns the partition index.
+// for sink it is always 0.
+func (s *UserDefinedSink) GetPartitionIdx() int32 {
+	return 0
+}
+
 func (s *UserDefinedSink) IsFull() bool {
 	return false
 }
diff --git a/pkg/sources/source.go b/pkg/sources/source.go
index 79fe7c55..09d20c40 100644
--- a/pkg/sources/source.go
+++ b/pkg/sources/source.go
@@ -85,13 +85,13 @@ func (sp *SourceProcessor) Start(ctx context.Context) error {
 				// TODO: remove this branch when deprecated limits are removed
 				writeOpts = append(writeOpts, redisclient.WithBufferUsageLimit(float64(*x.BufferUsageLimit)/100))
 			}
-			buffers := dfv1.GenerateBufferNames(sp.VertexInstance.Vertex.Namespace, sp.VertexInstance.Vertex.Spec.PipelineName, e.To, e.GetToVertexPartitionCount())
+			partitionedBuffers := dfv1.GenerateBufferNames(sp.VertexInstance.Vertex.Namespace, sp.VertexInstance.Vertex.Spec.PipelineName, e.To, e.GetToVertexPartitionCount())
 			var bufferWriters []isb.BufferWriter
 			// create a writer for each partition.
-			for _, partition := range buffers {
+			for partitionIdx, partition := range partitionedBuffers {
 				group := partition + "-group"
 				redisClient := redisclient.NewInClusterRedisClient()
-				writer := redisisb.NewBufferWrite(ctx, redisClient, partition, group, writeOpts...)
+				writer := redisisb.NewBufferWrite(ctx, redisClient, partition, group, int32(partitionIdx), writeOpts...)
 				bufferWriters = append(bufferWriters, writer)
 			}
 			writersMap[e.To] = bufferWriters
@@ -124,12 +124,12 @@ func (sp *SourceProcessor) Start(ctx context.Context) error {
 				writeOpts = append(writeOpts, jetstreamisb.WithBufferUsageLimit(float64(*x.BufferUsageLimit)/100))
 			}
 			var bufferWriters []isb.BufferWriter
-			buffers := dfv1.GenerateBufferNames(sp.VertexInstance.Vertex.Namespace, sp.VertexInstance.Vertex.Spec.PipelineName, e.To, e.GetToVertexPartitionCount())
+			partitionedBuffers := dfv1.GenerateBufferNames(sp.VertexInstance.Vertex.Namespace, sp.VertexInstance.Vertex.Spec.PipelineName, e.To, e.GetToVertexPartitionCount())
 			// create a writer for each partition.
-			for _, buffer := range buffers {
-				streamName := isbsvc.JetStreamName(buffer)
+			for partitionIdx, partition := range partitionedBuffers {
+				streamName := isbsvc.JetStreamName(partition)
 				jetStreamClient := jsclient.NewInClusterJetStreamClient()
-				writer, err := jetstreamisb.NewJetStreamBufferWriter(ctx, jetStreamClient, buffer, streamName, streamName, writeOpts...)
+				writer, err := jetstreamisb.NewJetStreamBufferWriter(ctx, jetStreamClient, partition, streamName, streamName, int32(partitionIdx), writeOpts...)
 				if err != nil {
 					return err
 				}
diff --git a/pkg/udf/common.go b/pkg/udf/common.go
index 684e555f..dcff34d0 100644
--- a/pkg/udf/common.go
+++ b/pkg/udf/common.go
@@ -88,9 +88,9 @@ func buildRedisBufferIO(ctx context.Context, vertexInstance *dfv1.VertexInstance
 			writeOpts = append(writeOpts, redisclient.WithBufferUsageLimit(float64(*x.BufferUsageLimit)/100))
 		}
 		var edgeBuffers []isb.BufferWriter
-		buffers := dfv1.GenerateBufferNames(vertexInstance.Vertex.Namespace, vertexInstance.Vertex.Spec.PipelineName, e.To, e.GetToVertexPartitionCount())
-		for _, buffer := range buffers {
-			writer := redisisb.NewBufferWrite(ctx, redisClient, buffer, buffer+"-group", writeOpts...)
+		partitionedBuffers := dfv1.GenerateBufferNames(vertexInstance.Vertex.Namespace, vertexInstance.Vertex.Spec.PipelineName, e.To, e.GetToVertexPartitionCount())
+		for partitionIdx, partition := range partitionedBuffers {
+			writer := redisisb.NewBufferWrite(ctx, redisClient, partition, partition+"-group", int32(partitionIdx), writeOpts...)
 			edgeBuffers = append(edgeBuffers, writer)
 		}
 		writers[e.To] = edgeBuffers
@@ -166,11 +166,11 @@ func buildJetStreamBufferIO(ctx context.Context, vertexInstance *dfv1.VertexInst
 			writeOpts = append(writeOpts, jetstreamisb.WithBufferUsageLimit(float64(*x.BufferUsageLimit)/100))
 		}
 
-		buffers := dfv1.GenerateBufferNames(vertexInstance.Vertex.Namespace, vertexInstance.Vertex.Spec.PipelineName, e.To, e.GetToVertexPartitionCount())
+		partitionedBuffers := dfv1.GenerateBufferNames(vertexInstance.Vertex.Namespace, vertexInstance.Vertex.Spec.PipelineName, e.To, e.GetToVertexPartitionCount())
 		var edgeBuffers []isb.BufferWriter
-		for _, buffer := range buffers {
-			streamName := isbsvc.JetStreamName(buffer)
-			writer, err := jetstreamisb.NewJetStreamBufferWriter(ctx, jsclient.NewInClusterJetStreamClient(), buffer, streamName, streamName, writeOpts...)
+		for partitionIdx, partition := range partitionedBuffers {
+			streamName := isbsvc.JetStreamName(partition)
+			writer, err := jetstreamisb.NewJetStreamBufferWriter(ctx, jsclient.NewInClusterJetStreamClient(), partition, streamName, streamName, int32(partitionIdx), writeOpts...)
 			if err != nil {
 				return nil, nil, err
 			}


Test patch:
diff --git a/pkg/forward/forward_test.go b/pkg/forward/forward_test.go
index d3dda294..4094370e 100644
--- a/pkg/forward/forward_test.go
+++ b/pkg/forward/forward_test.go
@@ -307,7 +307,6 @@ func TestNewInterStepDataForward(t *testing.T) {
 
 				}
 			}()
-			wg.Wait()
 
 			wg.Add(1)
 			go func() {
@@ -321,7 +320,7 @@ func TestNewInterStepDataForward(t *testing.T) {
 						break loop
 					default:
 						if len(otKeys2) == 0 {
-							otKeys2, _ = otStores["to1"].GetAllKeys(ctx)
+							otKeys2, _ = otStores["to2"].GetAllKeys(ctx)
 							time.Sleep(time.Millisecond * 10)
 						} else {
 							// NOTE: in this test we only have one processor to publish
@@ -393,36 +392,69 @@ func TestNewInterStepDataForward(t *testing.T) {
 			wg.Add(1)
 			go func() {
 				defer wg.Done()
-				otKeys1, _ := otStores["to1"].GetAllKeys(ctx)
-				for otKeys1 == nil {
-					otKeys1, _ = otStores["to1"].GetAllKeys(ctx)
-					time.Sleep(time.Millisecond * 10)
+				otKeys1, _ := otStores["to2"].GetAllKeys(ctx)
+			loop:
+				for {
+					select {
+					case <-ctx.Done():
+						assert.Fail(t, "context cancelled while waiting to get a valid watermark")
+						break loop
+					default:
+						if len(otKeys1) == 0 {
+							otKeys1, _ = otStores["to1"].GetAllKeys(ctx)
+							time.Sleep(time.Millisecond * 10)
+						} else {
+							// NOTE: in this test we only have one processor to publish
+							// so len(otKeys) should always be 1
+							otKeys1, _ = otStores["to1"].GetAllKeys(ctx)
+							otValue1, _ := otStores["to1"].GetValue(ctx, otKeys1[0])
+							otDecode1, _ := wmb.DecodeToWMB(otValue1)
+							// if the batch size is 1, then the watermark should be idle because
+							// one of partitions will be idle
+							if otDecode1.Idle {
+								break loop
+							} else {
+								time.Sleep(time.Millisecond * 10)
+							}
+						}
+					}
+
 				}
 			}()
-			wg.Wait()
-			// NOTE: in this test we only have one processor to publish
-			// so len(otKeys) should always be 1
-			otKeys1, _ := otStores["to1"].GetAllKeys(ctx)
-			otValue1, _ := otStores["to1"].GetValue(ctx, otKeys1[0])
-			otDecode1, _ := wmb.DecodeToWMB(otValue1)
-			assert.True(t, otDecode1.Idle)
 
 			wg.Add(1)
 			go func() {
 				defer wg.Done()
 				otKeys2, _ := otStores["to2"].GetAllKeys(ctx)
-				for otKeys2 == nil {
-					otKeys2, _ = otStores["to2"].GetAllKeys(ctx)
-					time.Sleep(time.Millisecond * 10)
+			loop:
+				for {
+					select {
+					case <-ctx.Done():
+						assert.Fail(t, "context cancelled while waiting to get a valid watermark")
+						break loop
+					default:
+						if len(otKeys2) == 0 {
+							otKeys2, _ = otStores["to2"].GetAllKeys(ctx)
+							time.Sleep(time.Millisecond * 10)
+						} else {
+							// NOTE: in this test we only have one processor to publish
+							// so len(otKeys) should always be 1
+							otKeys2, _ = otStores["to2"].GetAllKeys(ctx)
+							otValue2, _ := otStores["to2"].GetValue(ctx, otKeys2[0])
+							otDecode2, _ := wmb.DecodeToWMB(otValue2)
+							// if the batch size is 1, then the watermark should be idle because
+							// one of partitions will be idle
+							if otDecode2.Idle {
+								break loop
+							} else {
+								time.Sleep(time.Millisecond * 10)
+							}
+						}
+					}
+
 				}
 			}()
 			wg.Wait()
-			// NOTE: in this test we only have one processor to publish
-			// so len(otKeys) should always be 1
-			otKeys2, _ := otStores["to2"].GetAllKeys(ctx)
-			otValue2, _ := otStores["to2"].GetValue(ctx, otKeys2[0])
-			otDecode2, _ := wmb.DecodeToWMB(otValue2)
-			assert.True(t, otDecode2.Idle)
 
 			msgs := to11.GetMessages(1)
 			for len(msgs) == 0 || msgs[0].Kind != isb.WMB {
@@ -538,39 +570,68 @@ func TestNewInterStepDataForward(t *testing.T) {
 			go func() {
 				defer wg.Done()
 				otKeys1, _ := otStores["to1"].GetAllKeys(ctx)
-				for otKeys1 == nil {
-					otKeys1, _ = otStores["to1"].GetAllKeys(ctx)
-					time.Sleep(time.Millisecond * 10)
+			loop:
+				for {
+					select {
+					case <-ctx.Done():
+						assert.Fail(t, "context cancelled while waiting to get a valid watermark")
+						break loop
+					default:
+						if len(otKeys1) == 0 {
+							otKeys1, _ = otStores["to1"].GetAllKeys(ctx)
+							time.Sleep(time.Millisecond * 10)
+						} else {
+							// NOTE: in this test we only have one processor to publish
+							// so len(otKeys) should always be 1
+							otKeys1, _ = otStores["to1"].GetAllKeys(ctx)
+							otValue1, _ := otStores["to1"].GetValue(ctx, otKeys1[0])
+							otDecode1, _ := wmb.DecodeToWMB(otValue1)
+							if tt.batchSize > 1 && !otDecode1.Idle {
+								break loop
+							} else if tt.batchSize == 1 && otDecode1.Idle {
+								break loop
+							} else {
+								time.Sleep(time.Millisecond * 10)
+							}
+						}
+					}
+
 				}
 			}()
-			wg.Wait()
-			// NOTE: in this test we only have one processor to publish
-			// so len(otKeys) should always be 1
-			otKeys1, _ := otStores["to1"].GetAllKeys(ctx)
-			otValue1, _ := otStores["to1"].GetValue(ctx, otKeys1[0])
-			otDecode1, _ := wmb.DecodeToWMB(otValue1)
-			if tt.batchSize > 1 {
-				assert.False(t, otDecode1.Idle)
-			} else {
-				assert.True(t, otDecode1.Idle)
-			}
 
 			wg.Add(1)
 			go func() {
 				defer wg.Done()
 				otKeys2, _ := otStores["to2"].GetAllKeys(ctx)
-				for otKeys2 == nil {
-					otKeys2, _ = otStores["to2"].GetAllKeys(ctx)
-					time.Sleep(time.Millisecond * 10)
+			loop:
+				for {
+					select {
+					case <-ctx.Done():
+						assert.Fail(t, "context cancelled while waiting to get a valid watermark")
+						break loop
+					default:
+						if len(otKeys2) == 0 {
+							otKeys2, _ = otStores["to2"].GetAllKeys(ctx)
+							time.Sleep(time.Millisecond * 10)
+						} else {
+							// NOTE: in this test we only have one processor to publish
+							// so len(otKeys) should always be 1
+							otKeys2, _ = otStores["to2"].GetAllKeys(ctx)
+							otValue2, _ := otStores["to2"].GetValue(ctx, otKeys2[0])
+							otDecode2, _ := wmb.DecodeToWMB(otValue2)
+							// if the batch size is 1, then the watermark should be idle because
+							// one of partitions will be idle
+							if otDecode2.Idle {
+								break loop
+							} else {
+								time.Sleep(time.Millisecond * 10)
+							}
+						}
+					}
+
 				}
 			}()
 			wg.Wait()
-			// NOTE: in this test we only have one processor to publish
-			// so len(otKeys) should always be 1
-			otKeys2, _ := otStores["to2"].GetAllKeys(ctx)
-			otValue2, _ := otStores["to2"].GetValue(ctx, otKeys2[0])
-			otDecode2, _ := wmb.DecodeToWMB(otValue2)
-			assert.True(t, otDecode2.Idle)
 
 			// stop will cancel the contexts and therefore the forwarder stops without waiting
 			f.Stop()
diff --git a/pkg/isb/stores/jetstream/reader_test.go b/pkg/isb/stores/jetstream/reader_test.go
index f006d52f..7dfe4a60 100644
--- a/pkg/isb/stores/jetstream/reader_test.go
+++ b/pkg/isb/stores/jetstream/reader_test.go
@@ -36,6 +36,8 @@ func TestMain(m *testing.M) {
 	goleak.VerifyTestMain(m)
 }
 
+var defaultPartitionIdx = int32(0)
+
 func TestJetStreamBufferRead(t *testing.T) {
 	s := natstest.RunJetStreamServer(t)
 	defer natstest.ShutdownJetStreamServer(t, s)
@@ -54,7 +56,7 @@ func TestJetStreamBufferRead(t *testing.T) {
 	addStream(t, js, streamName)
 	defer deleteStream(js, streamName)
 
-	bw, err := NewJetStreamBufferWriter(ctx, defaultJetStreamClient, streamName, streamName, streamName)
+	bw, err := NewJetStreamBufferWriter(ctx, defaultJetStreamClient, streamName, streamName, streamName, defaultPartitionIdx)
 	assert.NoError(t, err)
 	jw, _ := bw.(*jetStreamWriter)
 	defer jw.Close()
@@ -81,7 +83,7 @@ func TestJetStreamBufferRead(t *testing.T) {
 		assert.NoError(t, e)
 	}
 
-	bufferReader, err := NewJetStreamBufferReader(ctx, defaultJetStreamClient, streamName, streamName, streamName, 0)
+	bufferReader, err := NewJetStreamBufferReader(ctx, defaultJetStreamClient, streamName, streamName, streamName, defaultPartitionIdx)
 	assert.NoError(t, err)
 
 	fromStep := bufferReader.(*jetStreamReader)
@@ -149,7 +151,7 @@ func TestGetName(t *testing.T) {
 	addStream(t, js, streamName)
 	defer deleteStream(js, streamName)
 
-	bufferReader, err := NewJetStreamBufferReader(ctx, defaultJetStreamClient, streamName, streamName, streamName, 0)
+	bufferReader, err := NewJetStreamBufferReader(ctx, defaultJetStreamClient, streamName, streamName, streamName, defaultPartitionIdx)
 	assert.NoError(t, err)
 	br := bufferReader.(*jetStreamReader)
 	assert.Equal(t, br.GetName(), streamName)
@@ -176,7 +178,7 @@ func TestClose(t *testing.T) {
 	addStream(t, js, streamName)
 	defer deleteStream(js, streamName)
 
-	bufferReader, err := NewJetStreamBufferReader(ctx, defaultJetStreamClient, streamName, streamName, streamName, 0)
+	bufferReader, err := NewJetStreamBufferReader(ctx, defaultJetStreamClient, streamName, streamName, streamName, defaultPartitionIdx)
 	assert.NoError(t, err)
 
 	br := bufferReader.(*jetStreamReader)
diff --git a/pkg/isb/stores/jetstream/writer_test.go b/pkg/isb/stores/jetstream/writer_test.go
index 9dc43d55..416b324d 100644
--- a/pkg/isb/stores/jetstream/writer_test.go
+++ b/pkg/isb/stores/jetstream/writer_test.go
@@ -90,7 +90,7 @@ func TestForwarderJetStreamBuffer(t *testing.T) {
 			addStream(t, js, toStreamName)
 			defer deleteStream(js, toStreamName)
 
-			bw, err := NewJetStreamBufferWriter(ctx, defaultJetStreamClient, streamName, streamName, streamName, WithMaxLength(10))
+			bw, err := NewJetStreamBufferWriter(ctx, defaultJetStreamClient, streamName, streamName, streamName, defaultPartitionIdx, WithMaxLength(10))
 			assert.NoError(t, err)
 			jw, _ := bw.(*jetStreamWriter)
 			defer jw.Close()
@@ -128,12 +128,12 @@ func TestForwarderJetStreamBuffer(t *testing.T) {
 			}}
 
 			// Forwarder logic tested here with a jetstream read and write
-			bufferReader, err := NewJetStreamBufferReader(ctx, defaultJetStreamClient, streamName, streamName, streamName, 0)
+			bufferReader, err := NewJetStreamBufferReader(ctx, defaultJetStreamClient, streamName, streamName, streamName, defaultPartitionIdx)
 			assert.NoError(t, err)
 			fromStep, _ := bufferReader.(*jetStreamReader)
 			defer fromStep.Close()
 
-			bufferWriter, err := NewJetStreamBufferWriter(ctx, defaultJetStreamClient, toStreamName, toStreamName, toStreamName, WithMaxLength(10))
+			bufferWriter, err := NewJetStreamBufferWriter(ctx, defaultJetStreamClient, toStreamName, toStreamName, toStreamName, defaultPartitionIdx, WithMaxLength(10))
 			assert.NoError(t, err)
 			to1 := bufferWriter.(*jetStreamWriter)
 			defer to1.Close()
@@ -204,7 +204,7 @@ func TestJetStreamBufferWriterBufferFull(t *testing.T) {
 	addStream(t, js, streamName)
 	defer deleteStream(js, streamName)
 
-	bw, err := NewJetStreamBufferWriter(ctx, defaultJetStreamClient, streamName, streamName, streamName, WithMaxLength(10), WithBufferUsageLimit(0.2))
+	bw, err := NewJetStreamBufferWriter(ctx, defaultJetStreamClient, streamName, streamName, streamName, defaultPartitionIdx, WithMaxLength(10), WithBufferUsageLimit(0.2))
 	assert.NoError(t, err)
 	jw, _ := bw.(*jetStreamWriter)
 	defer jw.Close()
@@ -263,7 +263,7 @@ func TestJetStreamBufferWriterBufferFull_DiscardLatest(t *testing.T) {
 	addStream(t, js, streamName)
 	defer deleteStream(js, streamName)
 
-	bw, err := NewJetStreamBufferWriter(ctx, defaultJetStreamClient, streamName, streamName, streamName, WithMaxLength(10), WithBufferUsageLimit(0.2), WithBufferFullWritingStrategy(dfv1.DiscardLatest))
+	bw, err := NewJetStreamBufferWriter(ctx, defaultJetStreamClient, streamName, streamName, streamName, defaultPartitionIdx, WithMaxLength(10), WithBufferUsageLimit(0.2), WithBufferFullWritingStrategy(dfv1.DiscardLatest))
 	assert.NoError(t, err)
 	jw, _ := bw.(*jetStreamWriter)
 	defer jw.Close()
@@ -321,7 +321,7 @@ func TestWriteGetName(t *testing.T) {
 	addStream(t, js, streamName)
 	defer deleteStream(js, streamName)
 
-	bufferWriter, err := NewJetStreamBufferReader(ctx, defaultJetStreamClient, streamName, streamName, streamName, 0)
+	bufferWriter, err := NewJetStreamBufferReader(ctx, defaultJetStreamClient, streamName, streamName, streamName, defaultPartitionIdx)
 	assert.NoError(t, err)
 
 	bw := bufferWriter.(*jetStreamReader)
@@ -349,7 +349,7 @@ func TestWriteClose(t *testing.T) {
 	addStream(t, js, streamName)
 	defer deleteStream(js, streamName)
 
-	bufferWriter, err := NewJetStreamBufferWriter(ctx, defaultJetStreamClient, streamName, streamName, streamName)
+	bufferWriter, err := NewJetStreamBufferWriter(ctx, defaultJetStreamClient, streamName, streamName, streamName, defaultPartitionIdx)
 	assert.NoError(t, err)
 
 	bw := bufferWriter.(*jetStreamWriter)
diff --git a/pkg/isb/stores/redis/read_test.go b/pkg/isb/stores/redis/read_test.go
index de2a0542..ccdad57b 100644
--- a/pkg/isb/stores/redis/read_test.go
+++ b/pkg/isb/stores/redis/read_test.go
@@ -48,6 +48,8 @@ var (
 	testStartTime = time.Unix(1636470000, 0).UTC()
 )
 
+var defaultPartitionIdx = int32(0)
+
 func TestRedisQRead_Read(t *testing.T) {
 	ctx := context.Background()
 	client := redisclient.NewRedisClient(redisOptions)
@@ -56,7 +58,7 @@ func TestRedisQRead_Read(t *testing.T) {
 	consumer := "con-0"
 
 	count := int64(10)
-	rqr, _ := NewBufferRead(ctx, client, stream, group, consumer, 0).(*BufferRead)
+	rqr, _ := NewBufferRead(ctx, client, stream, group, consumer, defaultPartitionIdx).(*BufferRead)
 	err := client.CreateStreamGroup(ctx, rqr.GetStreamName(), group, redisclient.ReadFromEarliest)
 	assert.NoError(t, err)
 
@@ -88,7 +90,7 @@ func TestRedisCheckBacklog(t *testing.T) {
 	consumer := "readbacklog-consumer"
 
 	count := int64(10)
-	rqr, _ := NewBufferRead(ctx, client, stream, group, consumer, 0).(*BufferRead)
+	rqr, _ := NewBufferRead(ctx, client, stream, group, consumer, defaultPartitionIdx).(*BufferRead)
 	err := client.CreateStreamGroup(ctx, rqr.GetStreamName(), group, redisclient.ReadFromEarliest)
 	assert.NoError(t, err)
 
@@ -121,7 +123,7 @@ func TestRedisCheckBacklog(t *testing.T) {
 		},
 	}}
 
-	rqw, _ := NewBufferWrite(ctx, client, "toStream", "toGroup", redisclient.WithInfoRefreshInterval(2*time.Millisecond), redisclient.WithLagDuration(time.Minute)).(*BufferWrite)
+	rqw, _ := NewBufferWrite(ctx, client, "toStream", "toGroup", defaultPartitionIdx, redisclient.WithInfoRefreshInterval(2*time.Millisecond), redisclient.WithLagDuration(time.Minute)).(*BufferWrite)
 	err = client.CreateStreamGroup(ctx, rqw.GetStreamName(), "toGroup", redisclient.ReadFromEarliest)
 	assert.NoError(t, err)
 
@@ -161,8 +163,8 @@ func (suite *ReadTestSuite) SetupSuite() {
 	group := "testsuitegroup1"
 	consumer := "testsuite-0"
 	count := int64(10)
-	rqw, _ := NewBufferWrite(ctx, client, stream, group).(*BufferWrite)
-	rqr, _ := NewBufferRead(ctx, client, stream, group, consumer, 0).(*BufferRead)
+	rqw, _ := NewBufferWrite(ctx, client, stream, group, defaultPartitionIdx).(*BufferWrite)
+	rqr, _ := NewBufferRead(ctx, client, stream, group, consumer, defaultPartitionIdx).(*BufferRead)
 
 	suite.ctx = ctx
 	suite.rclient = client
@@ -316,8 +318,8 @@ func (suite *ReadWritePerformance) SetupSuite() {
 	toGroup := "ReadWritePerformance-group-to"
 	consumer := "ReadWritePerformance-con-0"
 	count := int64(10000)
-	rqw, _ := NewBufferWrite(ctx, client, toStream, toGroup, redisclient.WithInfoRefreshInterval(2*time.Millisecond), redisclient.WithLagDuration(time.Minute), redisclient.WithMaxLength(20000)).(*BufferWrite)
-	rqr, _ := NewBufferRead(ctx, client, fromStream, fromGroup, consumer, 0).(*BufferRead)
+	rqw, _ := NewBufferWrite(ctx, client, toStream, toGroup, defaultPartitionIdx, redisclient.WithInfoRefreshInterval(2*time.Millisecond), redisclient.WithLagDuration(time.Minute), redisclient.WithMaxLength(20000)).(*BufferWrite)
+	rqr, _ := NewBufferRead(ctx, client, fromStream, fromGroup, consumer, defaultPartitionIdx).(*BufferRead)
 
 	toSteps := map[string][]isb.BufferWriter{
 		"to1": {rqw},
@@ -375,7 +377,7 @@ func TestReadWritePerformanceSuite(t *testing.T) {
 
 // TestReadWriteLatency is used to look at the latency during forward.
 func (suite *ReadWritePerformance) TestReadWriteLatency() {
-	_ = NewBufferRead(suite.ctx, suite.rclient, "ReadWritePerformance-to", "ReadWritePerformance-group-to", "consumer-0", 0)
+	_ = NewBufferRead(suite.ctx, suite.rclient, "ReadWritePerformance-to", "ReadWritePerformance-group-to", "consumer-0", defaultPartitionIdx)
 	suite.False(suite.rqw.IsFull())
 	var writeMessages = make([]isb.Message, 0, suite.count)
 
@@ -406,8 +408,8 @@ func (suite *ReadWritePerformance) TestReadWriteLatency() {
 
 // TestReadWriteLatencyPipelining performs the latency test during a forward.
 func (suite *ReadWritePerformance) TestReadWriteLatencyPipelining() {
-	suite.rqw, _ = NewBufferWrite(suite.ctx, suite.rclient, "ReadWritePerformance-to", "ReadWritePerformance-group-to", redisclient.WithInfoRefreshInterval(2*time.Second), redisclient.WithLagDuration(time.Minute), redisclient.WithoutPipelining(), redisclient.WithMaxLength(20000)).(*BufferWrite)
-	_ = NewBufferRead(suite.ctx, suite.rclient, "ReadWritePerformance-to", "ReadWritePerformance-group-to", "consumer-0", 0)
+	suite.rqw, _ = NewBufferWrite(suite.ctx, suite.rclient, "ReadWritePerformance-to", "ReadWritePerformance-group-to", defaultPartitionIdx, redisclient.WithInfoRefreshInterval(2*time.Second), redisclient.WithLagDuration(time.Minute), redisclient.WithoutPipelining(), redisclient.WithMaxLength(20000)).(*BufferWrite)
+	_ = NewBufferRead(suite.ctx, suite.rclient, "ReadWritePerformance-to", "ReadWritePerformance-group-to", "consumer-0", defaultPartitionIdx)
 
 	vertex := &dfv1.Vertex{Spec: dfv1.VertexSpec{
 		PipelineName: "testPipeline",
diff --git a/pkg/isb/stores/redis/write_test.go b/pkg/isb/stores/redis/write_test.go
index 472a5cf3..ca9d8237 100644
--- a/pkg/isb/stores/redis/write_test.go
+++ b/pkg/isb/stores/redis/write_test.go
@@ -40,7 +40,7 @@ func TestRedisQWrite_Write(t *testing.T) {
 	client := redisclient.NewRedisClient(redisOptions)
 	ctx, cancel := context.WithTimeout(context.Background(), time.Minute*10)
 	defer cancel()
-	rqw, _ := NewBufferWrite(ctx, client, "rediswrite", "test", redisclient.WithLagDuration(2*time.Millisecond), redisclient.WithInfoRefreshInterval(2*time.Millisecond)).(*BufferWrite)
+	rqw, _ := NewBufferWrite(ctx, client, "rediswrite", "test", defaultPartitionIdx, redisclient.WithLagDuration(2*time.Millisecond), redisclient.WithInfoRefreshInterval(2*time.Millisecond)).(*BufferWrite)
 
 	for rqw.IsFull() {
 		select {
@@ -91,13 +91,13 @@ func TestRedisQWrite_WithPipeline(t *testing.T) {
 	ctx := context.Background()
 	stream := "withPipeline"
 	count := int64(100)
-	rqw, _ := NewBufferWrite(ctx, client, stream, "test", redisclient.WithoutPipelining()).(*BufferWrite)
+	rqw, _ := NewBufferWrite(ctx, client, stream, "test", defaultPartitionIdx, redisclient.WithoutPipelining()).(*BufferWrite)
 	streamName := rqw.GetStreamName()
 	defer func() { _ = client.DeleteKeys(ctx, streamName) }()
 	group := "withPipeline-group"
 	consumer := "withPipeline-0"
 
-	rqr, _ := NewBufferRead(ctx, client, stream, group, consumer, 0).(*BufferRead)
+	rqr, _ := NewBufferRead(ctx, client, stream, group, consumer, defaultPartitionIdx).(*BufferRead)
 	err := client.CreateStreamGroup(ctx, rqr.GetStreamName(), group, redisclient.ReadFromEarliest)
 	assert.NoError(t, err)
 	defer func() { _ = client.DeleteStreamGroup(ctx, rqr.GetStreamName(), group) }()
@@ -117,14 +117,14 @@ func TestRedisQWrite_WithoutPipeline(t *testing.T) {
 	ctx := context.Background()
 	stream := "withoutPipeline"
 	count := int64(100)
-	rqw, _ := NewBufferWrite(ctx, client, stream, "test").(*BufferWrite)
+	rqw, _ := NewBufferWrite(ctx, client, stream, "test", defaultPartitionIdx).(*BufferWrite)
 
 	streamName := rqw.GetStreamName()
 	defer func() { _ = client.DeleteKeys(ctx, streamName) }()
 	group := "withoutPipeline-group"
 	consumer := "withoutPipeline-0"
 
-	rqr, _ := NewBufferRead(ctx, client, stream, group, consumer, 0).(*BufferRead)
+	rqr, _ := NewBufferRead(ctx, client, stream, group, consumer, defaultPartitionIdx).(*BufferRead)
 	err := client.CreateStreamGroup(ctx, rqr.GetStreamName(), group, redisclient.ReadFromEarliest)
 	assert.NoError(t, err)
 
@@ -144,7 +144,7 @@ func TestRedisQWrite_WithInfoRefreshInterval(t *testing.T) {
 	stream := "withInfoRefreshInterval"
 	count := int64(10)
 	group := "withInfoRefreshInterval-group"
-	rqw, _ := NewBufferWrite(ctx, client, stream, group, redisclient.WithInfoRefreshInterval(2*time.Millisecond), redisclient.WithLagDuration(2*time.Millisecond), redisclient.WithMaxLength(10)).(*BufferWrite)
+	rqw, _ := NewBufferWrite(ctx, client, stream, group, defaultPartitionIdx, redisclient.WithInfoRefreshInterval(2*time.Millisecond), redisclient.WithLagDuration(2*time.Millisecond), redisclient.WithMaxLength(10)).(*BufferWrite)
 	err := client.CreateStreamGroup(ctx, rqw.GetStreamName(), group, redisclient.ReadFromEarliest)
 	if err != nil {
 		t.Fatalf("error creating consumer group: %s", err)
@@ -184,7 +184,7 @@ func TestRedisQWrite_WithInfoRefreshInterval_WithBufferFullWritingStrategyIsDisc
 	stream := "withInfoRefreshInterval"
 	count := int64(10)
 	group := "withInfoRefreshInterval-group"
-	rqw, _ := NewBufferWrite(ctx, client, stream, group, redisclient.WithInfoRefreshInterval(2*time.Millisecond), redisclient.WithLagDuration(2*time.Millisecond), redisclient.WithMaxLength(10), redisclient.WithBufferFullWritingStrategy(dfv1.DiscardLatest)).(*BufferWrite)
+	rqw, _ := NewBufferWrite(ctx, client, stream, group, defaultPartitionIdx, redisclient.WithInfoRefreshInterval(2*time.Millisecond), redisclient.WithLagDuration(2*time.Millisecond), redisclient.WithMaxLength(10), redisclient.WithBufferFullWritingStrategy(dfv1.DiscardLatest)).(*BufferWrite)
 	err := client.CreateStreamGroup(ctx, rqw.GetStreamName(), group, redisclient.ReadFromEarliest)
 	if err != nil {
 		t.Fatalf("error creating consumer group: %s", err)
@@ -282,7 +282,7 @@ func Test_updateIsFullFlag(t *testing.T) {
 	group := "getConsumerLag-group"
 	count := int64(9)
 
-	rqw, _ := NewBufferWrite(ctx, client, stream, group, redisclient.WithInfoRefreshInterval(2*time.Millisecond), redisclient.WithMaxLength(10)).(*BufferWrite)
+	rqw, _ := NewBufferWrite(ctx, client, stream, group, defaultPartitionIdx, redisclient.WithInfoRefreshInterval(2*time.Millisecond), redisclient.WithMaxLength(10)).(*BufferWrite)
 	err := client.CreateStreamGroup(ctx, rqw.GetStreamName(), group, redisclient.ReadFromEarliest)
 	defer func() { _ = client.DeleteStreamGroup(ctx, rqw.GetStreamName(), group) }()
 
@@ -316,7 +316,7 @@ func Test_GetName(t *testing.T) {
 	stream := "getName"
 	group := "getName-group"
 	ctx := context.Background()
-	rqw, _ := NewBufferWrite(ctx, client, stream, group).(*BufferWrite)
+	rqw, _ := NewBufferWrite(ctx, client, stream, group, defaultPartitionIdx).(*BufferWrite)
 	assert.Equal(t, stream, rqw.GetName())
 
 }
@@ -326,7 +326,7 @@ func Test_GetLag(t *testing.T) {
 	stream := "getLag"
 	group := "getLag-group"
 	ctx := context.Background()
-	rqw, _ := NewBufferWrite(ctx, client, stream, group).(*BufferWrite)
+	rqw, _ := NewBufferWrite(ctx, client, stream, group, defaultPartitionIdx).(*BufferWrite)
 	assert.Equal(t, time.Duration(0), rqw.GetConsumerLag())
 }
 
@@ -334,7 +334,7 @@ func Test_GetRefreshFullError(t *testing.T) {
 	client := redisclient.NewRedisClient(redisOptions)
 	group := "getRefreshFull-group"
 	ctx := context.Background()
-	rqw, _ := NewBufferWrite(ctx, client, "", group).(*BufferWrite)
+	rqw, _ := NewBufferWrite(ctx, client, "", group, defaultPartitionIdx).(*BufferWrite)
 	expected := uint32(3)
 	assert.Equal(t, expected, rqw.GetRefreshFullError())
 }
@@ -372,17 +372,17 @@ func TestNewInterStepDataForwardRedis(t *testing.T) {
 	toStream := "toStep"
 
 	// fromStep, we also have to create a fromStepWrite here to write data to the from stream
-	fromStep, _ := NewBufferRead(ctx, client, fromStream, fromGroup, consumer, 0, redisclient.WithInfoRefreshInterval(2*time.Millisecond)).(*BufferRead)
+	fromStep, _ := NewBufferRead(ctx, client, fromStream, fromGroup, consumer, defaultPartitionIdx, redisclient.WithInfoRefreshInterval(2*time.Millisecond)).(*BufferRead)
 	_ = client.CreateStreamGroup(ctx, fromStep.GetStreamName(), fromGroup, redisclient.ReadFromEarliest)
-	fromStepWrite, _ := NewBufferWrite(ctx, client, fromStream, fromGroup, redisclient.WithInfoRefreshInterval(2*time.Second), redisclient.WithLagDuration(1*time.Millisecond)).(*BufferWrite)
+	fromStepWrite, _ := NewBufferWrite(ctx, client, fromStream, fromGroup, defaultPartitionIdx, redisclient.WithInfoRefreshInterval(2*time.Second), redisclient.WithLagDuration(1*time.Millisecond)).(*BufferWrite)
 	defer func() { _ = client.DeleteKeys(ctx, fromStepWrite.GetStreamName()) }()
 	defer func() { _ = client.DeleteKeys(ctx, fromStep.GetStreamName()) }()
 	defer func() { _ = client.DeleteStreamGroup(ctx, fromStep.GetStreamName(), fromGroup) }()
 
 	// toStep, we also have to create a toStepRead here to read from the toStep
-	to1Read, _ := NewBufferRead(ctx, client, toStream, toGroup, consumer, 0).(*BufferRead)
+	to1Read, _ := NewBufferRead(ctx, client, toStream, toGroup, consumer, defaultPartitionIdx).(*BufferRead)
 	_ = client.CreateStreamGroup(ctx, to1Read.GetStreamName(), toGroup, redisclient.ReadFromEarliest)
-	to1, _ := NewBufferWrite(ctx, client, toStream, toGroup, redisclient.WithLagDuration(1*time.Millisecond), redisclient.WithInfoRefreshInterval(2*time.Second), redisclient.WithMaxLength(17)).(*BufferWrite)
+	to1, _ := NewBufferWrite(ctx, client, toStream, toGroup, defaultPartitionIdx, redisclient.WithLagDuration(1*time.Millisecond), redisclient.WithInfoRefreshInterval(2*time.Second), redisclient.WithMaxLength(17)).(*BufferWrite)
 	defer func() { _ = client.DeleteKeys(ctx, to1.GetStreamName()) }()
 	defer func() { _ = client.DeleteKeys(ctx, to1.GetStreamName()) }()
 	defer func() { _ = client.DeleteStreamGroup(ctx, to1.GetStreamName(), toGroup) }()
@@ -424,12 +424,12 @@ func TestReadTimeout(t *testing.T) {
 	toStream := "to-st"
 
 	// fromStep, we also have to create a fromStepWrite here to write data to the from stream
-	fromStep, _ := NewBufferRead(ctx, client, fromStream, fromGroup, consumer, 0, redisclient.WithInfoRefreshInterval(2*time.Millisecond)).(*BufferRead)
+	fromStep, _ := NewBufferRead(ctx, client, fromStream, fromGroup, consumer, defaultPartitionIdx, redisclient.WithInfoRefreshInterval(2*time.Millisecond)).(*BufferRead)
 	_ = client.CreateStreamGroup(ctx, fromStep.GetStreamName(), fromGroup, redisclient.ReadFromEarliest)
 	defer func() { _ = client.DeleteKeys(ctx, fromStep.GetStreamName()) }()
 	defer func() { _ = client.DeleteStreamGroup(ctx, fromStep.GetStreamName(), fromGroup) }()
 
-	to1, _ := NewBufferWrite(ctx, client, toStream, toGroup, redisclient.WithLagDuration(1*time.Millisecond), redisclient.WithInfoRefreshInterval(2*time.Second)).(*BufferWrite)
+	to1, _ := NewBufferWrite(ctx, client, toStream, toGroup, defaultPartitionIdx, redisclient.WithLagDuration(1*time.Millisecond), redisclient.WithInfoRefreshInterval(2*time.Second)).(*BufferWrite)
 	_ = client.CreateStreamGroup(ctx, to1.GetStreamName(), toGroup, redisclient.ReadFromEarliest)
 
 	defer func() { _ = client.DeleteKeys(ctx, to1.GetStreamName()) }()
@@ -466,7 +466,7 @@ func TestXTrimOnIsFull(t *testing.T) {
 	group := "trim-group"
 	buffer := "trim"
 
-	rqw, _ := NewBufferWrite(ctx, client, buffer, group, redisclient.WithLagDuration(1*time.Millisecond), redisclient.WithInfoRefreshInterval(2*time.Millisecond), redisclient.WithMaxLength(10)).(*BufferWrite)
+	rqw, _ := NewBufferWrite(ctx, client, buffer, group, defaultPartitionIdx, redisclient.WithLagDuration(1*time.Millisecond), redisclient.WithInfoRefreshInterval(2*time.Millisecond), redisclient.WithMaxLength(10)).(*BufferWrite)
 	err := client.CreateStreamGroup(ctx, rqw.GetStreamName(), group, redisclient.ReadFromEarliest)
 	assert.NoError(t, err)
 
@@ -531,7 +531,7 @@ func TestSetWriteInfo(t *testing.T) {
 	group := "setWriteInfo-group"
 	buffer := "setWriteInfo"
 
-	rqw, _ := NewBufferWrite(ctx, client, buffer, group, redisclient.WithLagDuration(1*time.Millisecond), redisclient.WithInfoRefreshInterval(2*time.Millisecond), redisclient.WithRefreshBufferWriteInfo(false)).(*BufferWrite)
+	rqw, _ := NewBufferWrite(ctx, client, buffer, group, defaultPartitionIdx, redisclient.WithLagDuration(1*time.Millisecond), redisclient.WithInfoRefreshInterval(2*time.Millisecond), redisclient.WithRefreshBufferWriteInfo(false)).(*BufferWrite)
 	err := client.CreateStreamGroup(ctx, rqw.GetStreamName(), group, redisclient.ReadFromEarliest)
 	assert.NoError(t, err)
 
@@ -551,7 +551,7 @@ func TestSetWriteInfo(t *testing.T) {
 	}
 
 	// Read all the messages.
-	rqr, _ := NewBufferRead(ctx, client, buffer, group, "consumer", 0).(*BufferRead)
+	rqr, _ := NewBufferRead(ctx, client, buffer, group, "consumer", defaultPartitionIdx).(*BufferRead)
 
 	defer func() { _ = client.DeleteKeys(ctx, rqr.GetStreamName()) }()
 
diff --git a/pkg/sources/nats/nats_test.go b/pkg/sources/nats/nats_test.go
index 73a37b19..6468d457 100644
--- a/pkg/sources/nats/nats_test.go
+++ b/pkg/sources/nats/nats_test.go
@@ -83,6 +83,8 @@ func newInstance(t *testing.T, vi *dfv1.VertexInstance) (*natsSource, error) {
 }
 
 func Test_Single(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
 	server := natstest.RunNatsServer(t)
 	defer server.Shutdown()
 
@@ -105,7 +107,17 @@ func Test_Single(t *testing.T) {
 
 	msgs, err := ns.Read(context.Background(), 5)
 	assert.NoError(t, err)
-	assert.Equal(t, 3, len(msgs))
+	for len(msgs) != 3 {
+		select {
+		case <-ctx.Done():
+			t.Fatal("timeout waiting for messages")
+			return
+		default:
+			msgs, err = ns.Read(context.Background(), 5)
+			assert.NoError(t, err)
+			time.Sleep(10 * time.Millisecond)
+		}
+	}
 }
 
 func Test_Multiple(t *testing.T) {