Fix panic in teststream. (#36227)

Author: shunping

sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go

@@ -2056,7 +2056,7 @@ func (ss *stageState) String() string {
 	return fmt.Sprintf("[%v] IN: %v OUT: %v UP: %q %v, kind: %v", ss.ID, ss.input, ss.output, pcol, up, ss.kind)
 }
 
-// updateWatermarks performs the following operations:
+// updateWatermarks performs the following operations and returns a possible set of stages to refresh next or nil.
 //
 // Watermark_In'  = MAX(Watermark_In, MIN(U(TS_Pending), U(Watermark_InputPCollection)))
 // Watermark_Out' = MAX(Watermark_Out, MIN(Watermark_In', U(minWatermarkHold)))

sdks/go/pkg/beam/runners/prism/internal/engine/teststream.go

@@ -253,6 +253,9 @@ func (ev tsFinalEvent) Execute(em *ElementManager) {
 	em.testStreamHandler.UpdateHold(em, mtime.MaxTimestamp)
 	ss := em.stages[ev.stageID]
 	kickSet := ss.updateWatermarks(em)
+	if kickSet == nil {
+		kickSet = make(set[string])
+	}
 	kickSet.insert(ev.stageID)
 	em.changedStages.merge(kickSet)
 }