Repository: numaproj/numaflow
Pull request #239 (2022-10-19 09:07:20)

test: reduce test with inmemory watermark store
Signed-off-by: Yashash H L <yashashhl25@gmail.com>

reduce tests integrated with in-memory watermark store

<!--

Before you push your changes:

* Run `make pre-push -B` to fix codegen and lint problems.

Then, you MUST:

* Sign-off your commit (otherwise the DCO check will fail).
* Use [a conventional commit message](https://www.conventionalcommits.org/en/v1.0.0/) (otherwise the commit message check will fail).

If you did not do this, reset all your commit and replace them with a single commit:

```
git reset HEAD~1 ;# change 1 to how many commits you made
git commit --signoff -m 'feat: my feat. Fixes #1234'
```

When creating your PR: 

* Make sure that "Fixes #" or "Closes #" is in both the PR title (for release notes) and description (to automatically link and close the issue).
* Say how you tested your changes. If you changed the UI, attach screenshots.
* Set your PR as a draft initially.
* Your PR needs to pass the required checks before it can be approved. 
* Once required tests have passed, mark your PR "Ready for review".

If changes were requested, once you've made them, you MUST dismiss the review to get it reviewed again.

-->



Patch:
diff --git a/pkg/isb/stores/simplebuffer/buffer.go b/pkg/isb/stores/simplebuffer/buffer.go
index 08076b45..96d1b293 100644
--- a/pkg/isb/stores/simplebuffer/buffer.go
+++ b/pkg/isb/stores/simplebuffer/buffer.go
@@ -95,6 +95,7 @@ func (b *InMemoryBuffer) IsEmpty() bool {
 
 func (b *InMemoryBuffer) Write(_ context.Context, messages []isb.Message) ([]isb.Offset, []error) {
 	var errs = make([]error, len(messages))
+	writeOffsets := make([]isb.Offset, len(messages))
 	for idx, message := range messages {
 		if !b.IsFull() {
 			var err1 error
@@ -111,14 +112,16 @@ func (b *InMemoryBuffer) Write(_ context.Context, messages []isb.Message) ([]isb
 			errs[idx] = nil
 			b.buffer[currentIdx].dirty = true
 			b.writeIdx = (currentIdx + 1) % b.size
-
+			writeOffsets = append(writeOffsets, isb.SimpleOffset(func() string {
+				return fmt.Sprintf("%d", currentIdx)
+			}))
 			// access buffer via lock
 			b.rwlock.Unlock()
 		} else {
 			errs[idx] = isb.BufferWriteErr{Name: b.name, Full: true, Message: "Buffer full!"}
 		}
 	}
-	return nil, errs
+	return writeOffsets, errs
 }
 
 func (b *InMemoryBuffer) blockIfEmpty(ctx context.Context) error {


Test patch:
diff --git a/pkg/reduce/pnf/processandforward_test.go b/pkg/reduce/pnf/processandforward_test.go
index 1559abd2..99eb8ac4 100644
--- a/pkg/reduce/pnf/processandforward_test.go
+++ b/pkg/reduce/pnf/processandforward_test.go
@@ -3,6 +3,7 @@ package pnf
 import (
 	"context"
 	"encoding/json"
+	"math"
 	"strings"
 	"testing"
 	"time"
@@ -10,6 +11,9 @@ import (
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/stores/simplebuffer"
 	"github.com/numaproj/numaflow/pkg/watermark/generic"
+	"github.com/numaproj/numaflow/pkg/watermark/processor"
+	"github.com/numaproj/numaflow/pkg/watermark/publish"
+	"github.com/numaproj/numaflow/pkg/watermark/store/inmem"
 
 	"github.com/golang/mock/gomock"
 	functionpb "github.com/numaproj/numaflow-go/pkg/apis/proto/function/v1"
@@ -21,6 +25,7 @@ import (
 	"github.com/numaproj/numaflow/pkg/pbq/partition"
 	"github.com/numaproj/numaflow/pkg/pbq/store"
 	udfcall "github.com/numaproj/numaflow/pkg/udf/function"
+	wmstore "github.com/numaproj/numaflow/pkg/watermark/store"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -153,11 +158,12 @@ func TestProcessAndForward_Forward(t *testing.T) {
 	}
 
 	tests := []struct {
-		name     string
-		id       partition.ID
-		buffers  []*simplebuffer.InMemoryBuffer
-		pf       ProcessAndForward
-		expected []bool
+		name       string
+		id         partition.ID
+		buffers    []*simplebuffer.InMemoryBuffer
+		pf         ProcessAndForward
+		expected   []bool
+		wmExpected map[string]int64
 	}{
 		{
 			name: "test-forward-one",
@@ -169,6 +175,10 @@ func TestProcessAndForward_Forward(t *testing.T) {
 			buffers:  []*simplebuffer.InMemoryBuffer{test1Buffer1, test1Buffer2},
 			pf:       createProcessAndForward(ctx, "test-forward-one", pbqManager, toBuffers1),
 			expected: []bool{false, true},
+			wmExpected: map[string]int64{
+				"buffer1": 120,
+				"buffer2": math.MinInt64,
+			},
 		},
 		{
 			name: "test-forward-all",
@@ -180,6 +190,10 @@ func TestProcessAndForward_Forward(t *testing.T) {
 			buffers:  []*simplebuffer.InMemoryBuffer{test2Buffer1, test2Buffer2},
 			pf:       createProcessAndForward(ctx, "test-forward-all", pbqManager, toBuffers2),
 			expected: []bool{false, false},
+			wmExpected: map[string]int64{
+				"buffer1": 120,
+				"buffer2": 120,
+			},
 		},
 		{
 			name: "test-drop-all",
@@ -191,6 +205,10 @@ func TestProcessAndForward_Forward(t *testing.T) {
 			buffers:  []*simplebuffer.InMemoryBuffer{test3Buffer1, test3Buffer2},
 			pf:       createProcessAndForward(ctx, "test-drop-all", pbqManager, toBuffers3),
 			expected: []bool{true, true},
+			wmExpected: map[string]int64{
+				"buffer1": math.MinInt64,
+				"buffer2": math.MinInt64,
+			},
 		},
 	}
 
@@ -199,7 +217,14 @@ func TestProcessAndForward_Forward(t *testing.T) {
 			err := value.pf.Forward(ctx)
 			assert.NoError(t, err)
 			assert.Equal(t, []bool{value.buffers[0].IsEmpty(), value.buffers[1].IsEmpty()}, value.expected)
+			// pbq entry from the manager will be removed after forwarding
 			assert.Equal(t, pbqManager.GetPBQ(value.id), nil)
+			index := 0
+			for k, v := range value.pf.publishWatermark {
+				// expected watermark should be equal to window end time
+				assert.Equal(t, v.GetLatestWatermark().Unix(), value.wmExpected[k])
+				index += 1
+			}
 		})
 	}
 }
@@ -213,7 +238,7 @@ func createProcessAndForward(ctx context.Context, key string, pbqManager *pbq.Ma
 	}
 
 	// create a pbq for a partition
-	_, pw := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toBuffers)
+	pw := buildPublisherMap(toBuffers)
 	var simplePbq pbq.Reader
 	simplePbq, _ = pbqManager.CreateNewPBQ(ctx, testPartition)
 
@@ -247,3 +272,23 @@ func createProcessAndForward(ctx context.Context, key string, pbqManager *pbq.Ma
 
 	return pf
 }
+
+// buildPublisherMap builds publisher for each toBuffer
+func buildPublisherMap(toBuffers map[string]isb.BufferWriter) map[string]publish.Publisher {
+	var ctx = context.Background()
+
+	var publisherHBKeyspace = "publisherTest_PROCESSORS"
+
+	var publisherOTKeyspace = "publisherTest_OT_publisherTestPod1"
+
+	heartbeatKV, _, _ := inmem.NewKVInMemKVStore(ctx, "testPublisher", publisherHBKeyspace)
+	otKV, _, _ := inmem.NewKVInMemKVStore(ctx, "testPublisher", publisherOTKeyspace)
+
+	publishers := make(map[string]publish.Publisher)
+	for key := range toBuffers {
+		publishEntity := processor.NewProcessorEntity(key)
+		p := publish.NewPublish(ctx, publishEntity, wmstore.BuildWatermarkStore(heartbeatKV, otKV), publish.WithAutoRefreshHeartbeatDisabled(), publish.WithPodHeartbeatRate(1))
+		publishers[key] = p
+	}
+	return publishers
+}
diff --git a/pkg/reduce/reduce_test.go b/pkg/reduce/reduce_test.go
index 0dc3ae44..e5246a63 100644
--- a/pkg/reduce/reduce_test.go
+++ b/pkg/reduce/reduce_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"math"
 	"strconv"
 	"testing"
 	"time"
@@ -13,8 +14,11 @@ import (
 	"github.com/numaproj/numaflow/pkg/isb/stores/simplebuffer"
 	"github.com/numaproj/numaflow/pkg/pbq"
 	"github.com/numaproj/numaflow/pkg/pbq/store"
+	"github.com/numaproj/numaflow/pkg/watermark/fetch"
 	"github.com/numaproj/numaflow/pkg/watermark/processor"
 	"github.com/numaproj/numaflow/pkg/watermark/publish"
+	wmstore "github.com/numaproj/numaflow/pkg/watermark/store"
+	"github.com/numaproj/numaflow/pkg/watermark/store/inmem"
 	"github.com/numaproj/numaflow/pkg/window/strategy/fixed"
 	"github.com/stretchr/testify/assert"
 )
@@ -46,11 +50,11 @@ type PayloadForTest struct {
 	Value int
 }
 
-type forwardReduceTest struct {
+type CounterReduceTest struct {
 }
 
-// Reduce returns a result, with last message's event time
-func (f forwardReduceTest) Reduce(ctx context.Context, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
+// Reduce returns a result with the count of messages
+func (f CounterReduceTest) Reduce(ctx context.Context, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
 	count := 0
 	for range messageStream {
 		count += 1
@@ -73,14 +77,72 @@ func (f forwardReduceTest) Reduce(ctx context.Context, messageStream <-chan *isb
 	}, nil
 }
 
-func (f forwardReduceTest) WhereTo(s string) ([]string, error) {
+func (f CounterReduceTest) WhereTo(s string) ([]string, error) {
 	return []string{"to"}, nil
 }
 
+type SumReduceTest struct {
+}
+
+func (s SumReduceTest) Reduce(ctx context.Context, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
+	sum := 0
+	for msg := range messageStream {
+		var payload PayloadForTest
+		_ = json.Unmarshal(msg.Payload, &payload)
+		sum += payload.Value
+	}
+
+	payload := PayloadForTest{Key: "sum", Value: sum}
+	b, _ := json.Marshal(payload)
+	ret := &isb.Message{
+		Header: isb.Header{
+			PaneInfo: isb.PaneInfo{
+				EventTime: time.Now(),
+			},
+			ID:  "msgID",
+			Key: "result",
+		},
+		Body: isb.Body{Payload: b},
+	}
+	return []*isb.Message{
+		ret,
+	}, nil
+}
+
+type MaxReduceTest struct {
+}
+
+func (m MaxReduceTest) Reduce(ctx context.Context, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
+	mx := math.MinInt64
+	for msg := range messageStream {
+		var payload PayloadForTest
+		_ = json.Unmarshal(msg.Payload, &payload)
+		if payload.Value > mx {
+			mx = payload.Value
+		}
+	}
+
+	payload := PayloadForTest{Key: "max", Value: mx}
+	b, _ := json.Marshal(payload)
+	ret := &isb.Message{
+		Header: isb.Header{
+			PaneInfo: isb.PaneInfo{
+				EventTime: time.Now(),
+			},
+			ID:  "msgID",
+			Key: "result",
+		},
+		Body: isb.Body{Payload: b},
+	}
+	return []*isb.Message{
+		ret,
+	}, nil
+}
+
 // read from simple buffer
 // mock reduce op to return result
 // assert to check if the result is forwarded to toBuffers
-func TestDataForward_Start(t *testing.T) {
+func TestDataForward_StartWithNoOpWM(t *testing.T) {
 	windowTime := 2 * time.Second
 	parentCtx := context.Background()
 	child, cancelFn := context.WithTimeout(parentCtx, windowTime*2)
@@ -91,7 +153,7 @@ func TestDataForward_Start(t *testing.T) {
 	to := simplebuffer.NewInMemoryBuffer("to", toBufferSize)
 
 	// keep on writing <count> messages every 1 second for the supplied key
-	go writeMessages(child, 10, "test-1", fromBuffer)
+	go writeMessages(child, 10, "no-op-test", fromBuffer, EventTypeWMProgressor{}, time.Second*1)
 
 	toBuffer := map[string]isb.BufferWriter{
 		"to": to,
@@ -113,7 +175,7 @@ func TestDataForward_Start(t *testing.T) {
 	window := fixed.NewFixed(windowTime)
 
 	var reduceDataForwarder *DataForward
-	reduceDataForwarder, err = NewDataForward(child, forwardReduceTest{}, fromBuffer, toBuffer, pbqManager, forwardReduceTest{}, EventTypeWMProgressor{}, publisher, window, WithReadBatchSize(10))
+	reduceDataForwarder, err = NewDataForward(child, CounterReduceTest{}, fromBuffer, toBuffer, pbqManager, CounterReduceTest{}, EventTypeWMProgressor{}, publisher, window, WithReadBatchSize(10))
 	assert.NoError(t, err)
 
 	go reduceDataForwarder.Start(child)
@@ -138,38 +200,193 @@ func TestDataForward_Start(t *testing.T) {
 	var readMessagePayload PayloadForTest
 	_ = json.Unmarshal(msgs[0].Payload, &readMessagePayload)
 	// here the count will always be equal to the number of messages written per key
-	assert.Equal(t, 10, readMessagePayload.Value)
+	assert.Equal(t, 1, readMessagePayload.Value)
 	assert.Equal(t, "count", readMessagePayload.Key)
 
 }
 
-func writeMessages(ctx context.Context, count int, key string, fromBuffer *simplebuffer.InMemoryBuffer) {
-	generateTime := 1 * time.Second
-	ticker := time.NewTicker(generateTime)
-	i := 1
-	for {
-		select {
-		case <-ticker.C:
-			// build 10 messages from the start time, time difference between the messages is 1 second
-			messages := buildMessagesForReduce(count, key+strconv.Itoa(i))
-			i++
+func TestDataForward_StartWithInMemoryWMStore(t *testing.T) {
+	var err error
+	fromBufferSize := int64(100000)
+	toBufferSize := int64(10)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+
+	//create from buffers for tests
+	fromBuffer1 := simplebuffer.NewInMemoryBuffer("from1", fromBufferSize)
+	fromBuffer2 := simplebuffer.NewInMemoryBuffer("from2", fromBufferSize)
+	fromBuffer3 := simplebuffer.NewInMemoryBuffer("from3", fromBufferSize)
+
+	//create to buffers for tests
+	buffer1 := simplebuffer.NewInMemoryBuffer("to", toBufferSize)
+	buffer2 := simplebuffer.NewInMemoryBuffer("to", toBufferSize)
+	buffer3 := simplebuffer.NewInMemoryBuffer("to", toBufferSize)
+	toBuffer1 := map[string]isb.BufferWriter{
+		"to": buffer1,
+	}
+	toBuffer2 := map[string]isb.BufferWriter{
+		"to": buffer2,
+	}
+	toBuffer3 := map[string]isb.BufferWriter{
+		"to": buffer3,
+	}
 
-			// TODO use watermark in memory store when it is ready to synthesize offsets
-			// write the messages to fromBuffer, so that it will be available for consuming
-			fromBuffer.Write(ctx, messages)
-		case <-ctx.Done():
-			ticker.Stop()
-			return
+	// create pbqManager for tests
+	var pbqManager1, pbqManager2, pbqManager3 *pbq.Manager
+	pbqManager1, err = pbq.NewManager(ctx, pbq.WithPBQStoreOptions(store.WithStoreSize(100), store.WithPbqStoreType(dfv1.InMemoryType)),
+		pbq.WithReadTimeout(1*time.Second), pbq.WithChannelBufferSize(10))
+	assert.NoError(t, err)
+
+	pbqManager2, err = pbq.NewManager(ctx, pbq.WithPBQStoreOptions(store.WithStoreSize(100), store.WithPbqStoreType(dfv1.InMemoryType)),
+		pbq.WithReadTimeout(1*time.Second), pbq.WithChannelBufferSize(10))
+	assert.NoError(t, err)
+
+	pbqManager3, err = pbq.NewManager(ctx, pbq.WithPBQStoreOptions(store.WithStoreSize(10000), store.WithPbqStoreType(dfv1.InMemoryType)),
+		pbq.WithReadTimeout(1*time.Second), pbq.WithChannelBufferSize(10))
+	assert.NoError(t, err)
+
+	// create watermark fetcher and publisher for tests
+	f1, p1 := fetcherAndPublisher(ctx, toBuffer1, fromBuffer1, "test-1")
+	f2, p2 := fetcherAndPublisher(ctx, toBuffer2, fromBuffer2, "test-2")
+	f3, p3 := fetcherAndPublisher(ctx, toBuffer3, fromBuffer3, "test-3")
+
+	// create go routine sources which produces messages
+	// keep on writing <count> messages every 1 second for the supplied key
+	go writeMessages(ctx, 10, "test-1", fromBuffer1, p1["from1"], time.Second*2)
+	go writeMessages(ctx, 100, "test-2", fromBuffer2, p2["from2"], time.Minute*1)
+	go writeMessages(ctx, 1000, "test-3", fromBuffer3, p3["from3"], time.Minute*10)
+
+	//create window for tests
+	window1 := fixed.NewFixed(2 * time.Second)
+	window2 := fixed.NewFixed(2 * time.Minute)
+	window3 := fixed.NewFixed(20 * time.Minute)
+
+	//create forwarder for tests
+	var reduceDataForwarder1, reduceDataForwarder2, reduceDataForwarder3 *DataForward
+	reduceDataForwarder1, err = NewDataForward(ctx, CounterReduceTest{}, fromBuffer1, toBuffer1, pbqManager1, CounterReduceTest{}, f1, p1, window1, WithReadBatchSize(10))
+	assert.NoError(t, err)
+	reduceDataForwarder2, err = NewDataForward(ctx, SumReduceTest{}, fromBuffer2, toBuffer2, pbqManager2, CounterReduceTest{}, f2, p2, window2, WithReadBatchSize(10))
+	assert.NoError(t, err)
+	reduceDataForwarder3, err = NewDataForward(ctx, MaxReduceTest{}, fromBuffer3, toBuffer3, pbqManager3, CounterReduceTest{}, f3, p3, window3, WithReadBatchSize(10))
+	assert.NoError(t, err)
+
+	tests := []struct {
+		name                string
+		window              *fixed.Fixed
+		fromBuffer          isb.BufferReader
+		toBuffer            *simplebuffer.InMemoryBuffer
+		reduceDataForwarder *DataForward
+		expectedKey         string
+		expectedValue       int64
+	}{
+		{
+			name:                "test-1",
+			window:              window1,
+			fromBuffer:          fromBuffer1,
+			toBuffer:            buffer1,
+			reduceDataForwarder: reduceDataForwarder1,
+			expectedKey:         "count",
+			expectedValue:       10,
+		},
+		{
+			name:                "test-2",
+			window:              window2,
+			fromBuffer:          fromBuffer2,
+			toBuffer:            buffer2,
+			reduceDataForwarder: reduceDataForwarder2,
+			expectedKey:         "sum",
+			expectedValue:       4950,
+		},
+		{
+			name:                "test-3",
+			window:              window3,
+			fromBuffer:          fromBuffer3,
+			toBuffer:            buffer3,
+			reduceDataForwarder: reduceDataForwarder3,
+			expectedKey:         "max",
+			expectedValue:       999,
+		},
+	}
+
+	for _, value := range tests {
+		t.Run(value.name, func(t *testing.T) {
+			go value.reduceDataForwarder.Start(ctx)
+
+			// wait until there is data in to buffer
+			for value.toBuffer.IsEmpty() {
+				select {
+				case <-ctx.Done():
+				default:
+					time.Sleep(100 * time.Millisecond)
+				}
+			}
+
+			// we are reading only one message here but the count should be equal to
+			// the number of keyed windows that closed
+			msgs, readErr := value.toBuffer.Read(ctx, 1)
+			assert.Nil(t, readErr)
+			assert.Len(t, msgs, 1)
+
+			// assert the output of reduce
+			var readMessagePayload PayloadForTest
+			_ = json.Unmarshal(msgs[0].Payload, &readMessagePayload)
+			assert.Equal(t, value.expectedValue, int64(readMessagePayload.Value))
+			assert.Equal(t, value.expectedKey, readMessagePayload.Key)
+		})
+	}
+
+}
+
+// fetcherAndPublisher creates watermark fetcher and publishers for toBuffers, and keeps the processors alive by sending heartbeats
+func fetcherAndPublisher(ctx context.Context, toBuffers map[string]isb.BufferWriter, fromBuffer *simplebuffer.InMemoryBuffer, key string) (fetch.Fetcher, map[string]publish.Publisher) {
+
+	var (
+		keyspace     = key
+		pipelineName = "testPipeline"
+		hbBucketName = keyspace + "_PROCESSORS"
+		otBucketName = keyspace + "_OT"
+	)
+
+	publishEntity := processor.NewProcessorEntity("reduceProcessor" + keyspace)
+	sourcePublishEntity := processor.NewProcessorEntity("sourceProcessor" + keyspace)
+	hb, hbWatcherCh, _ := inmem.NewKVInMemKVStore(ctx, pipelineName, hbBucketName)
+	ot, otWatcherCh, _ := inmem.NewKVInMemKVStore(ctx, pipelineName, otBucketName)
+
+	publishers := make(map[string]publish.Publisher)
+
+	p := publish.NewPublish(ctx, publishEntity, wmstore.BuildWatermarkStore(hb, ot), publish.WithAutoRefreshHeartbeatDisabled(), publish.WithPodHeartbeatRate(1))
+	sourcePublisher := publish.NewPublish(ctx, sourcePublishEntity, wmstore.BuildWatermarkStore(hb, ot), publish.WithAutoRefreshHeartbeatDisabled(), publish.WithPodHeartbeatRate(1))
+
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			default:
+				_ = hb.PutKV(ctx, "sourceProcessor"+keyspace, []byte(fmt.Sprintf("%d", time.Now().Unix())))
+				_ = hb.PutKV(ctx, "reduceProcessor"+keyspace, []byte(fmt.Sprintf("%d", time.Now().Unix())))
+				time.Sleep(time.Duration(1) * time.Second)
+			}
 		}
+	}()
+
+	publishers[fromBuffer.GetName()] = sourcePublisher
+	for key := range toBuffers {
+		publishers[key] = p
 	}
 
+	hbWatcher, _ := inmem.NewInMemWatch(ctx, pipelineName, keyspace+"_PROCESSORS", hbWatcherCh)
+	otWatcher, _ := inmem.NewInMemWatch(ctx, pipelineName, keyspace+"_OT", otWatcherCh)
+
+	var pm = fetch.NewProcessorManager(ctx, wmstore.BuildWatermarkStoreWatcher(hbWatcher, otWatcher), fetch.WithPodHeartbeatRate(1), fetch.WithRefreshingProcessorsRate(1), fetch.WithSeparateOTBuckets(false))
+	var f = fetch.NewEdgeFetcher(ctx, "from", pm)
+	return f, publishers
 }
 
 // buildMessagesForReduce builds test isb.Message which can be used for testing reduce.
-func buildMessagesForReduce(count int, key string) []isb.Message {
+func buildMessagesForReduce(count int, key string, publishTime time.Time) []isb.Message {
 	var messages = make([]isb.Message, 0, count)
 	for i := 0; i < count; i++ {
-		tmpTime := time.Now()
 		result, _ := json.Marshal(PayloadForTest{
 			Key:   key,
 			Value: i,
@@ -178,7 +395,7 @@ func buildMessagesForReduce(count int, key string) []isb.Message {
 			isb.Message{
 				Header: isb.Header{
 					PaneInfo: isb.PaneInfo{
-						EventTime: tmpTime,
+						EventTime: publishTime,
 					},
 					ID:  fmt.Sprintf("%d", i),
 					Key: key,
@@ -190,3 +407,30 @@ func buildMessagesForReduce(count int, key string) []isb.Message {
 
 	return messages
 }
+
+// writeMessages writes message to simple buffer and publishes the watermark using write offsets
+func writeMessages(ctx context.Context, count int, key string, fromBuffer *simplebuffer.InMemoryBuffer, publish publish.Publisher, interval time.Duration) {
+	generateTime := 100 * time.Millisecond
+	ticker := time.NewTicker(generateTime)
+	intervalDuration := interval
+	publishTime := time.Unix(60, 0)
+	i := 1
+	for {
+		select {
+		case <-ticker.C:
+			// build  messages with eventTime time.Now()
+			publishTime = publishTime.Add(intervalDuration)
+			messages := buildMessagesForReduce(count, key+strconv.Itoa(i), publishTime)
+			i++
+
+			// write the messages to fromBuffer, so that it will be available for consuming
+			offsets, _ := fromBuffer.Write(ctx, messages)
+			if len(offsets) > 0 {
+				publish.PublishWatermark(processor.Watermark(publishTime), offsets[len(offsets)-1])
+			}
+		case <-ctx.Done():
+			ticker.Stop()
+			return
+		}
+	}
+}
diff --git a/pkg/watermark/publish/publisher_inmem_test.go b/pkg/watermark/publish/publisher_inmem_test.go
index c9978959..5d0a4f42 100644
--- a/pkg/watermark/publish/publisher_inmem_test.go
+++ b/pkg/watermark/publish/publisher_inmem_test.go
@@ -15,7 +15,7 @@ import (
 	"github.com/numaproj/numaflow/pkg/watermark/store"
 )
 
-func TestPublisherWithSeparateOTBuckets_InMem(t *testing.T) {
+func TestPublisherWithSharedOTBuckets_InMem(t *testing.T) {
 	var ctx = context.Background()
 
 	var publisherHBKeyspace = "publisherTest_PROCESSORS"
@@ -58,7 +58,7 @@ func TestPublisherWithSeparateOTBuckets_InMem(t *testing.T) {
 
 }
 
-func TestPublisherWithSharedOTBucket_InMem(t *testing.T) {
+func TestPublisherWithSeparateOTBucket_InMem(t *testing.T) {
 	var ctx = context.Background()
 
 	var keyspace = "publisherTest"