3 files changed, 83 insertions, 8 deletions

diff --git a/core/operator/harness_test.go b/core/operator/harness_test.go
new file mode 100644
index 0000000..383ab8c
--- /dev/null
+++ b/core/operator/harness_test.go
@@ -0,0 +1,23 @@
+package operator_test
+
+import (
+        "context"
+
+        "flux/core"
+)
+
+type captureEmitter[T any] struct {
+        records []core.StreamRecord[T]
+}
+
+func (e *captureEmitter[T]) Emit(_ context.Context, record core.StreamRecord[T]) error {
+        e.records = append(e.records, record)
+        return nil
+}
+
+func (e *captureEmitter[T]) Records() []core.StreamRecord[T] {
+        result := make([]core.StreamRecord[T], len(e.records))
+        copy(result, e.records)
+
+        return result
+}
diff --git a/core/operator/watermark.go b/core/operator/watermark.go
index 3d591e5..43d63fa 100644
--- a/core/operator/watermark.go
+++ b/core/operator/watermark.go
@@ -6,20 +6,15 @@ import (
         "flux/core"
 )
 
-// TimestampAssigner is responsible for event time defined by user
-type TimestampAssigner[IN any] func(in IN) int64
-
 type WaterGeneratorOperator[IN any] struct {
-        assigner             TimestampAssigner[IN]
         maxOutOfOrder        int64
         currentMaxTimestamp  int64
         lastEmittedWatermark int64
 }
 
-func NewWatermarkGeneratorOperator[IN any](maxOutOfOrder int64, assigner TimestampAssigner[IN]) *WaterGeneratorOperator[IN] {
+func NewWatermarkGeneratorOperator[IN any](maxOutOfOrder int64) *WaterGeneratorOperator[IN] {
         return &WaterGeneratorOperator[IN]{
                 maxOutOfOrder:        maxOutOfOrder,
-                assigner:             assigner,
                 lastEmittedWatermark: -1,
         }
 }
@@ -34,8 +29,7 @@ func (w *WaterGeneratorOperator[IN]) Process(ctx context.Context, record core.St
                         return err
                 }
                 // query the time
-                eventTime := w.assigner(record.Value)
-                record.Timestamp = eventTime
+                eventTime := record.Timestamp
 
                 // update the max timestamp
                 if w.currentMaxTimestamp < eventTime {
diff --git a/core/operator/watermark_test.go b/core/operator/watermark_test.go
new file mode 100644
index 0000000..0a0495e
--- /dev/null
+++ b/core/operator/watermark_test.go
@@ -0,0 +1,58 @@
+package operator_test
+
+import (
+        "context"
+        "flux/core"
+        "flux/core/operator"
+        "testing"
+)
+
+type timestampEvent struct {
+        EventTime int64
+}
+
+func TestWatermarkGeneratorPreservesAssignedEventTimestamp(t *testing.T) {
+        ctx := context.Background()
+
+        op := operator.NewWatermarkGeneratorOperator[timestampEvent](0)
+
+        if err := op.Open(ctx); err != nil {
+                t.Fatalf("open operator failed: %v", err)
+        }
+
+        t.Cleanup(func() {
+                if err := op.Close(ctx); err != nil {
+                        t.Errorf("close operator: %v", err)
+                }
+        })
+
+        out := &captureEmitter[timestampEvent]{}
+        input := core.NewDataRecord("user-1", timestampEvent{EventTime: 999}, 40)
+
+        if err := op.Process(ctx, input, out); err != nil {
+                t.Fatalf("process operator failed: %v", err)
+        }
+
+        records := out.Records()
+        if len(records) == 0 {
+                t.Fatal("operator emitted no records")
+        }
+
+        got := records[0]
+
+        if got.Type != core.RecordData {
+                t.Fatalf("first output type: got %v, want RecordData", got.Type)
+        }
+
+        if got.Key != "user-1" {
+                t.Errorf("data key: got %q, want %q", got.Key, "user-1")
+        }
+
+        if got.Value != input.Value {
+                t.Errorf("data value: got %+v, want %+v", got.Value, input.Value)
+        }
+
+        if got.Timestamp != 40 {
+                t.Errorf("data timestamp: got %d, want %d", got.Timestamp, 40)
+        }
+}