[BEAM-14546] Fix errant pass for empty collections in Count

sdks/go/pkg/beam/io/synthetic/step.go

@@ -49,9 +49,9 @@ func init() {
 func Step(s beam.Scope, cfg StepConfig, col beam.PCollection) beam.PCollection {
 	s = s.Scope("synthetic.Step")
 	if cfg.Splittable {
-		return beam.ParDo(s, &sdfStepFn{cfg: cfg}, col)
+		return beam.ParDo(s, &sdfStepFn{Cfg: cfg}, col)
 	}
-	return beam.ParDo(s, &stepFn{cfg: cfg}, col)
+	return beam.ParDo(s, &stepFn{Cfg: cfg}, col)
 }
 
 // stepFn is a DoFn implementing behavior for synthetic steps. For usage
@@ -60,7 +60,7 @@ func Step(s beam.Scope, cfg StepConfig, col beam.PCollection) beam.PCollection {
 // The stepFn is expected to be initialized with a cfg and will follow that
 // config to determine its behavior when emitting elements.
 type stepFn struct {
-	cfg StepConfig
+	Cfg StepConfig
 	rng randWrapper
 }
 
@@ -73,9 +73,9 @@ func (fn *stepFn) Setup() {
 // outputs identical to that input based on the outputs per input configuration
 // in StepConfig.
 func (fn *stepFn) ProcessElement(key, val []byte, emit func([]byte, []byte)) {
-	filtered := fn.cfg.FilterRatio > 0 && fn.rng.Float64() < fn.cfg.FilterRatio
+	filtered := fn.Cfg.FilterRatio > 0 && fn.rng.Float64() < fn.Cfg.FilterRatio
 
-	for i := 0; i < fn.cfg.OutputPerInput; i++ {
+	for i := 0; i < fn.Cfg.OutputPerInput; i++ {
 		if !filtered {
 			emit(key, val)
 		}
@@ -88,7 +88,7 @@ func (fn *stepFn) ProcessElement(key, val []byte, emit func([]byte, []byte)) {
 // The sdfStepFn is expected to be initialized with a cfg and will follow
 // that config to determine its behavior when splitting and emitting elements.
 type sdfStepFn struct {
-	cfg StepConfig
+	Cfg StepConfig
 	rng randWrapper
 }
 
@@ -98,7 +98,7 @@ type sdfStepFn struct {
 func (fn *sdfStepFn) CreateInitialRestriction(_, _ []byte) offsetrange.Restriction {
 	return offsetrange.Restriction{
 		Start: 0,
-		End:   int64(fn.cfg.OutputPerInput),
+		End:   int64(fn.Cfg.OutputPerInput),
 	}
 }
 
@@ -107,7 +107,7 @@ func (fn *sdfStepFn) CreateInitialRestriction(_, _ []byte) offsetrange.Restricti
 // method will contain at least one element, so the number of splits will not
 // exceed the number of elements.
 func (fn *sdfStepFn) SplitRestriction(_, _ []byte, rest offsetrange.Restriction) (splits []offsetrange.Restriction) {
-	return rest.EvenSplits(int64(fn.cfg.InitialSplits))
+	return rest.EvenSplits(int64(fn.Cfg.InitialSplits))
 }
 
 // RestrictionSize outputs the size of the restriction as the number of elements
@@ -130,7 +130,7 @@ func (fn *sdfStepFn) Setup() {
 // ProcessElement takes an input and either filters it or produces a number of
 // outputs identical to that input based on the restriction size.
 func (fn *sdfStepFn) ProcessElement(rt *sdf.LockRTracker, key, val []byte, emit func([]byte, []byte)) {
-	filtered := fn.cfg.FilterRatio > 0 && fn.rng.Float64() < fn.cfg.FilterRatio
+	filtered := fn.Cfg.FilterRatio > 0 && fn.rng.Float64() < fn.Cfg.FilterRatio
 
 	for i := rt.GetRestriction().(offsetrange.Restriction).Start; rt.TryClaim(i); i++ {
 		if !filtered {

sdks/go/pkg/beam/io/synthetic/step_test.go

@@ -37,7 +37,7 @@ func TestStepConfig_OutputPerInput(t *testing.T) {
 			cfg := DefaultStepConfig().OutputPerInput(test.outPer).Build()
 
 			// Non-splittable StepFn.
-			dfn := stepFn{cfg: cfg}
+			dfn := stepFn{Cfg: cfg}
 			var keys [][]byte
 			emitFn := func(key []byte, val []byte) {
 				keys = append(keys, key)
@@ -52,7 +52,7 @@ func TestStepConfig_OutputPerInput(t *testing.T) {
 
 			// SDF StepFn.
 			cfg = DefaultStepConfig().OutputPerInput(test.outPer).Splittable(true).Build()
-			sdf := sdfStepFn{cfg: cfg}
+			sdf := sdfStepFn{Cfg: cfg}
 			keys, _ = simulateSdfStepFn(t, &sdf)
 			if got := len(keys); got != test.outPer {
 				t.Errorf("sdfStepFn emitted wrong number of outputs: got: %v, want: %v",
@@ -97,7 +97,7 @@ func TestStepConfig_FilterRatio(t *testing.T) {
 
 			// Non-splittable StepFn.
 			cfg := DefaultStepConfig().FilterRatio(test.ratio).Build()
-			dfn := stepFn{cfg: cfg}
+			dfn := stepFn{Cfg: cfg}
 			dfn.Setup()
 			dfn.rng = &fakeRand{f64: test.rand}
 			dfn.ProcessElement(elm, elm, emitFn)
@@ -110,7 +110,7 @@ func TestStepConfig_FilterRatio(t *testing.T) {
 
 			// SDF StepFn.
 			cfg = DefaultStepConfig().FilterRatio(test.ratio).Splittable(true).Build()
-			sdf := sdfStepFn{cfg: cfg}
+			sdf := sdfStepFn{Cfg: cfg}
 			keys = nil
 			rest := sdf.CreateInitialRestriction(elm, elm)
 			splits := sdf.SplitRestriction(elm, elm, rest)
@@ -154,7 +154,7 @@ func TestStepConfig_InitialSplits(t *testing.T) {
 					Build()
 				elm := []byte{0, 0, 0, 0}
 
-				sdf := sdfStepFn{cfg: cfg}
+				sdf := sdfStepFn{Cfg: cfg}
 				rest := sdf.CreateInitialRestriction(elm, elm)
 				splits := sdf.SplitRestriction(elm, elm, rest)
 				if got := len(splits); got != test.want {
@@ -186,7 +186,7 @@ func TestStepConfig_InitialSplits(t *testing.T) {
 					InitialSplits(test.splits).
 					Build()
 
-				sdf := sdfStepFn{cfg: cfg}
+				sdf := sdfStepFn{Cfg: cfg}
 				keys, _ := simulateSdfStepFn(t, &sdf)
 				if got := len(keys); got != test.want {
 					t.Errorf("SourceFn emitted wrong number of outputs: got: %v, want: %v",

sdks/go/pkg/beam/testing/passert/count.go

@@ -30,6 +30,10 @@ func Count(s beam.Scope, col beam.PCollection, name string, count int) {
 	if typex.IsKV(col.Type()) {
 		col = beam.DropKey(s, col)
 	}
+
+	if count > 0 {
+		NonEmpty(s, col)
+	}
 	counted := beam.Combine(s, &elmCountCombineFn{}, col)
 	beam.ParDo0(s, &errFn{Name: name, Count: count}, counted)
 }

sdks/go/pkg/beam/testing/passert/count_test.go

@@ -22,24 +22,62 @@ import (
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/ptest"
 )
 
-func TestCount_Good(t *testing.T) {
-	p, s := beam.NewPipelineWithRoot()
-	col := beam.Create(s, "a", "b", "c", "d", "e")
-	count := 5
+func TestCount(t *testing.T) {
+	var tests = []struct {
+		name     string
+		elements []string
+		count    int
+	}{
+		{
+			"full",
+			[]string{"a", "b", "c", "d", "e"},
+			5,
+		},
+		{
+			"empty",
+			[]string{},
+			0,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			p, s := beam.NewPipelineWithRoot()
+			col := beam.CreateList(s, test.elements)
 
-	Count(s, col, "TestCount_Good", count)
-	if err := ptest.Run(p); err != nil {
-		t.Errorf("Pipeline failed: %v", err)
+			Count(s, col, test.name, test.count)
+			if err := ptest.Run(p); err != nil {
+				t.Errorf("Pipeline failed: %v", err)
+			}
+		})
 	}
 }
 
 func TestCount_Bad(t *testing.T) {
-	p, s := beam.NewPipelineWithRoot()
-	col := beam.Create(s, "a", "b", "c", "d", "e")
-	count := 10
+	var tests = []struct {
+		name     string
+		elements []string
+		count    int
+	}{
+		{
+			"mismatch",
+			[]string{"a", "b", "c", "d", "e"},
+			10,
+		},
+		{
+			"empty pcollection",
+			[]string{},
+			5,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			p, s := beam.NewPipelineWithRoot()
+			col := beam.CreateList(s, test.elements)
 
-	Count(s, col, "TestCount_Bad", count)
-	if err := ptest.Run(p); err == nil {
-		t.Errorf("pipeline SUCCEEDED but should have failed")
+			Count(s, col, test.name, test.count)
+			if err := ptest.Run(p); err == nil {
+				t.Errorf("pipeline SUCCEEDED but should have failed")
+			}
+		})
 	}
 }

sdks/go/pkg/beam/testing/passert/hash.go

@@ -31,6 +31,7 @@ import (
 func Hash(s beam.Scope, col beam.PCollection, name, hash string, size int) {
 	s = s.Scope(fmt.Sprintf("passert.Hash(%v)", name))
 
+	NonEmpty(s, col)
 	keyed := beam.AddFixedKey(s, col)
 	grouped := beam.GroupByKey(s, keyed)
 	beam.ParDo0(s, &hashFn{Name: name, Size: size, Hash: hash}, grouped)

sdks/go/pkg/beam/testing/passert/hash_test.go

@@ -0,0 +1,33 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package passert
+
+import (
+	"testing"
+
+	"github.com/apache/beam/sdks/v2/go/pkg/beam"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/ptest"
+)
+
+func TestHash_Bad(t *testing.T) {
+	p, s := beam.NewPipelineWithRoot()
+	col := beam.CreateList(s, []string{})
+
+	Hash(s, col, "empty collection", "", 0)
+	if err := ptest.Run(p); err == nil {
+		t.Errorf("pipeline SUCCEEDED but should have failed")
+	}
+}

sdks/go/pkg/beam/testing/passert/passert.go

@@ -29,7 +29,7 @@ import (
 )
 
 //go:generate go install github.com/apache/beam/sdks/v2/go/cmd/starcgen
-//go:generate starcgen --package=passert --identifiers=diffFn,failFn,failIfBadEntries,failKVFn,failGBKFn,hashFn,sumFn,errFn,elmCountCombineFn
+//go:generate starcgen --package=passert --identifiers=diffFn,failFn,failIfBadEntries,failKVFn,failGBKFn,hashFn,sumFn,errFn,elmCountCombineFn,nonEmptyFn
 //go:generate go fmt
 
 // Diff splits 2 incoming PCollections into 3: left only, both, right only. Duplicates are
@@ -179,3 +179,20 @@ type failGBKFn struct {
 func (f *failGBKFn) ProcessElement(x beam.X, _ func(*beam.Y) bool) error {
 	return errors.Errorf(f.Format, fmt.Sprintf("(%v,*)", x))
 }
+
+type nonEmptyFn struct{}
+
+func (n *nonEmptyFn) ProcessElement(_ []byte, iter func(*beam.Z) bool) error {
+	var val beam.Z
+	for iter(&val) {
+		return nil
+	}
+	return errors.New("PCollection is empty, want non-empty collection")
+}
+
+// NonEmpty asserts that the given PCollection has at least one element.
+func NonEmpty(s beam.Scope, col beam.PCollection) beam.PCollection {
+	s = s.Scope("passert.NonEmpty")
+	beam.ParDo0(s, &nonEmptyFn{}, beam.Impulse(s), beam.SideInput{Input: col})
+	return col
+}