Set state integration test (#22935)

* WIP: Set state integration test

* Update sdks/go/test/integration/primitives/state.go

Co-authored-by: Ritesh Ghorse <riteshghorse@gmail.com>

Co-authored-by: Ritesh Ghorse <riteshghorse@gmail.com>

sdks/go/test/integration/integration.go

@@ -97,6 +97,7 @@ var directFilters = []string{
  "TestBagState_Clear",
  "TestCombiningState",
  "TestMapState",
+ "TestSetState",
 }
 
 var portableFilters = []string{
@@ -125,6 +126,7 @@ var portableFilters = []string{
  "TestBagState_Clear",
  "TestCombiningState",
  "TestMapState",
+ "TestSetState",
 }
 
 var flinkFilters = []string{
@@ -174,6 +176,7 @@ var samzaFilters = []string{
  "TestBagState_Clear",
  "TestCombiningState",
  "TestMapState",
+ "TestSetState",
 }
 
 var sparkFilters = []string{

sdks/go/test/integration/primitives/state.go

@@ -36,6 +36,7 @@ func init() {
  register.DoFn3x1[state.Provider, string, int, string](&bagStateClearFn{})
  register.DoFn3x1[state.Provider, string, int, string](&combiningStateFn{})
  register.DoFn3x1[state.Provider, string, int, string](&mapStateFn{})
+ register.DoFn3x1[state.Provider, string, int, string](&setStateFn{})
  register.Emitter2[string, int]()
  register.Combiner1[int](&combine1{})
  register.Combiner2[string, int](&combine2{})
@@ -400,3 +401,47 @@ func MapStateParDo() *beam.Pipeline {
 
  return p
 }
+
+type setStateFn struct {
+ State1 state.Set[string]
+}
+
+func (f *setStateFn) ProcessElement(s state.Provider, w string, c int) string {
+ ok, err := f.State1.Contains(s, w)
+ if err != nil {
+ panic(err)
+ }
+ err = f.State1.Add(s, w)
+ if err != nil {
+ panic(err)
+ }
+ if ok {
+ err = f.State1.Add(s, fmt.Sprintf("%v%v", w, 1))
+ if err != nil {
+ panic(err)
+ }
+ }
+
+ keys, _, err := f.State1.Keys(s)
+ if err != nil {
+ panic(err)
+ }
+
+ sort.Slice(keys, func(i, j int) bool { return keys[i] < keys[j] })
+
+ return fmt.Sprintf("%v: %v, keys: %v", w, ok, keys)
+}
+
+// SetStateParDo tests a DoFn that uses set state.
+func SetStateParDo() *beam.Pipeline {
+ p, s := beam.NewPipelineWithRoot()
+
+ in := beam.Create(s, "apple", "pear", "peach", "apple", "apple", "pear")
+ keyed := beam.ParDo(s, func(w string, emit func(string, int)) {
+ emit(w, 1)
+ }, in)
+ counts := beam.ParDo(s, &setStateFn{State1: state.MakeSetState[string]("key1")}, keyed)
+ passert.Equals(s, counts, "apple: false, keys: [apple]", "pear: false, keys: [pear]", "peach: false, keys: [peach]", "apple: true, keys: [apple apple1]", "apple: true, keys: [apple apple1]", "pear: true, keys: [pear pear1]")
+
+ return p
+}

sdks/go/test/integration/primitives/state_test.go

@@ -56,3 +56,8 @@ func TestMapState(t *testing.T) {
  integration.CheckFilters(t)
  ptest.RunAndValidate(t, MapStateParDo())
 }
+
+func TestSetState(t *testing.T) {
+ integration.CheckFilters(t)
+ ptest.RunAndValidate(t, SetStateParDo())
+}