[GitHub] [beam] youngoli commented on a change in pull request #11763: [BEAM-9978] Adding functionality and tests to Go offset range tracker.
youngoli commented on a change in pull request #11763:
URL: https://github.com/apache/beam/pull/11763#discussion_r429026079
##
File path: sdks/go/pkg/beam/io/rtrackers/offsetrange/offsetrange_test.go
##
@@ -0,0 +1,212 @@
+// 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 offsetrange
+
+import (
+ "fmt"
+ "github.com/google/go-cmp/cmp"
+ "testing"
+)
+
+// TestRestriction_EvenSplits tests various splits and checks that they all
+// follow the contract for EvenSplits. This means that all restrictions are
+// evenly split, that each restriction has at least one element, and that each
+// element is present in the split restrictions.
+func TestRestriction_EvenSplits(t *testing.T) {
+ tests := []struct {
+ rest Restriction
+ num int64
+ }{
+ {rest: Restriction{Start: 0, End: 21}, num: 4},
+ {rest: Restriction{Start: 21, End: 42}, num: 4},
+ {rest: Restriction{Start: 0, End: 5}, num: 10},
+ {rest: Restriction{Start: 0, End: 21}, num: -1},
+ }
+ for _, test := range tests {
+ test := test
+ t.Run(fmt.Sprintf("(rest[%v, %v], splits = %v)",
+ test.rest.Start, test.rest.End, test.num), func(t
*testing.T) {
+ r := test.rest
+
+ // Get the minimum size that a split restriction can
be. Max size
+ // should be min + 1. This way we can check the size of
each split.
+ num := test.num
+ if num <= 1 {
+ num = 1
+ }
+ min := (r.End - r.Start) / num
+
+ splits := r.EvenSplits(test.num)
+ prevEnd := r.Start
+ for _, split := range splits {
+ size := split.End - split.Start
+ // Check: Each restriction has at least 1
element.
+ if size == 0 {
+ t.Errorf("split restriction [%v, %v] is
empty, size must be greater than 0.",
+ split.Start, split.End)
+ }
+ // Check: Restrictions are evenly split.
+ if size != min && size != min+1 {
+ t.Errorf("split restriction [%v, %v]
has unexpected size. got: %v, want: %v or %v",
+ split.Start, split.End, size,
min, min+1)
+ }
+ // Check: All elements are still in a split
restrictions. This
+ // logic assumes that the splits are returned
in order which
+ // isn't guaranteed by EvenSplits, but this
check is way easier
+ // with the assumption.
+ if split.Start != prevEnd {
+ t.Errorf("restriction range [%v, %v]
missing after splits.",
+ prevEnd, split.Start)
+ } else {
+ prevEnd = split.End
+ }
+ }
+ if prevEnd != r.End {
+ t.Errorf("restriction range [%v, %v] missing
after splits.",
+ prevEnd, r.End)
+ }
+ })
+ }
+}
+
+// TestTracker_TryClaim validates both success and failure cases for TryClaim.
+func TestTracker_TryClaim(t *testing.T) {
+ // Test that TryClaim works as expected when called correctly.
+ t.Run("Correctness", func(t *testing.T) {
+ tests := []struct {
+ rest Restriction
+ claims []int64
+ }{
+ {rest: Restriction{Start: 0, End: 3}, claims:
[]int64{0, 1, 2, 3}},
+ {rest: Restriction{Start: 10, End:
[GitHub] [beam] youngoli commented on a change in pull request #11763: [BEAM-9978] Adding functionality and tests to Go offset range tracker.
youngoli commented on a change in pull request #11763:
URL: https://github.com/apache/beam/pull/11763#discussion_r429025895
##
File path: sdks/go/pkg/beam/io/rtrackers/offsetrange/offsetrange_test.go
##
@@ -0,0 +1,212 @@
+// 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 offsetrange
+
+import (
+ "fmt"
+ "github.com/google/go-cmp/cmp"
+ "testing"
+)
+
+// TestRestriction_EvenSplits tests various splits and checks that they all
+// follow the contract for EvenSplits. This means that all restrictions are
+// evenly split, that each restriction has at least one element, and that each
+// element is present in the split restrictions.
+func TestRestriction_EvenSplits(t *testing.T) {
+ tests := []struct {
+ rest Restriction
+ num int64
+ }{
+ {rest: Restriction{Start: 0, End: 21}, num: 4},
+ {rest: Restriction{Start: 21, End: 42}, num: 4},
+ {rest: Restriction{Start: 0, End: 5}, num: 10},
+ {rest: Restriction{Start: 0, End: 21}, num: -1},
+ }
+ for _, test := range tests {
+ test := test
+ t.Run(fmt.Sprintf("(rest[%v, %v], splits = %v)",
+ test.rest.Start, test.rest.End, test.num), func(t
*testing.T) {
+ r := test.rest
+
+ // Get the minimum size that a split restriction can
be. Max size
+ // should be min + 1. This way we can check the size of
each split.
+ num := test.num
+ if num <= 1 {
+ num = 1
+ }
+ min := (r.End - r.Start) / num
+
+ splits := r.EvenSplits(test.num)
+ prevEnd := r.Start
+ for _, split := range splits {
+ size := split.End - split.Start
+ // Check: Each restriction has at least 1
element.
+ if size == 0 {
+ t.Errorf("split restriction [%v, %v] is
empty, size must be greater than 0.",
+ split.Start, split.End)
+ }
+ // Check: Restrictions are evenly split.
+ if size != min && size != min+1 {
+ t.Errorf("split restriction [%v, %v]
has unexpected size. got: %v, want: %v or %v",
+ split.Start, split.End, size,
min, min+1)
+ }
+ // Check: All elements are still in a split
restrictions. This
+ // logic assumes that the splits are returned
in order which
+ // isn't guaranteed by EvenSplits, but this
check is way easier
+ // with the assumption.
+ if split.Start != prevEnd {
+ t.Errorf("restriction range [%v, %v]
missing after splits.",
+ prevEnd, split.Start)
+ } else {
+ prevEnd = split.End
+ }
+ }
+ if prevEnd != r.End {
+ t.Errorf("restriction range [%v, %v] missing
after splits.",
+ prevEnd, r.End)
+ }
+ })
+ }
+}
+
+// TestTracker_TryClaim validates both success and failure cases for TryClaim.
+func TestTracker_TryClaim(t *testing.T) {
+ // Test that TryClaim works as expected when called correctly.
+ t.Run("Correctness", func(t *testing.T) {
+ tests := []struct {
+ rest Restriction
+ claims []int64
+ }{
+ {rest: Restriction{Start: 0, End: 3}, claims:
[]int64{0, 1, 2, 3}},
+ {rest: Restriction{Start: 10, End:
