github.com/apache/beam/sdks/v2@v2.48.2/go/test/regression/lperror_test.go (about)

     1  // Licensed to the Apache Software Foundation (ASF) under one or more
     2  // contributor license agreements.  See the NOTICE file distributed with
     3  // this work for additional information regarding copyright ownership.
     4  // The ASF licenses this file to You under the Apache License, Version 2.0
     5  // (the "License"); you may not use this file except in compliance with
     6  // the License.  You may obtain a copy of the License at
     7  //
     8  //    http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  
    16  package regression
    17  
    18  import (
    19  	"testing"
    20  
    21  	"github.com/apache/beam/sdks/v2/go/pkg/beam"
    22  	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/metrics"
    23  	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/passert"
    24  	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/ptest"
    25  	"github.com/apache/beam/sdks/v2/go/test/integration"
    26  
    27  	_ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/dataflow"
    28  	_ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/flink"
    29  	_ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/samza"
    30  	_ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/spark"
    31  )
    32  
    33  func TestLPErrorPipeline(t *testing.T) {
    34  	integration.CheckFilters(t)
    35  
    36  	pipeline, s := beam.NewPipelineWithRoot()
    37  	want := beam.CreateList(s, []int{0})
    38  	got := LPErrorPipeline(s)
    39  	passert.Equals(s, got, want)
    40  
    41  	ptest.RunAndValidate(t, pipeline)
    42  }
    43  
    44  func TestLPErrorReshufflePipeline(t *testing.T) {
    45  	integration.CheckFilters(t)
    46  
    47  	pipeline, s := beam.NewPipelineWithRoot()
    48  	LPErrorReshufflePipeline(s)
    49  
    50  	ptest.RunAndValidate(t, pipeline)
    51  }
    52  
    53  func TestLPErrorReshufflePipeline_PAssert(t *testing.T) {
    54  	integration.CheckFilters(t)
    55  
    56  	pipeline, s := beam.NewPipelineWithRoot()
    57  	got := LPErrorReshufflePipeline(s)
    58  	passert.Equals(s, got, fruit{"Apple"}, fruit{"Banana"}, fruit{"Cherry"})
    59  
    60  	ptest.RunAndValidate(t, pipeline)
    61  }
    62  
    63  func TestLPErrorReshufflePipeline_SideInput(t *testing.T) {
    64  	integration.CheckFilters(t)
    65  
    66  	pipeline, s := beam.NewPipelineWithRoot()
    67  	got := LPErrorReshufflePipeline(s)
    68  	beam.ParDo0(s, &iterSideStrings{
    69  		Wants: []string{"Apple", "Banana", "Cherry"},
    70  	}, beam.Impulse(s), beam.SideInput{Input: got})
    71  
    72  	ptest.RunAndValidate(t, pipeline)
    73  }
    74  
    75  func checkFruitCount(t *testing.T, pr beam.PipelineResult) {
    76  	t.Helper()
    77  	fcr := pr.Metrics().Query(func(sr metrics.SingleResult) bool {
    78  		return sr.Namespace() == MetricNamespace && sr.Name() == FruitCounterName
    79  	})
    80  	if len(fcr.Counters()) == 0 {
    81  		t.Logf("no counters found: check if %v supports counters", *ptest.Runner)
    82  		return
    83  	}
    84  	if got, want := fcr.Counters()[0].Result(), int64(3); got != want {
    85  		t.Errorf("unexpected fruit count: got %v, want %v", got, want)
    86  	}
    87  }
    88  
    89  func TestLPErrorReshufflePipeline_DoFn(t *testing.T) {
    90  	integration.CheckFilters(t)
    91  
    92  	pipeline, s := beam.NewPipelineWithRoot()
    93  	got := LPErrorReshufflePipeline(s)
    94  	beam.ParDo(s, countFruit, got)
    95  
    96  	pr := ptest.RunAndValidate(t, pipeline)
    97  	checkFruitCount(t, pr)
    98  }
    99  
   100  func TestLPErrorReshufflePipeline_DoFnPAssert(t *testing.T) {
   101  	integration.CheckFilters(t)
   102  
   103  	pipeline, s := beam.NewPipelineWithRoot()
   104  	got := LPErrorReshufflePipeline(s)
   105  	counted := beam.ParDo(s, countFruit, got)
   106  	passert.Equals(s, counted, fruit{"Apple"}, fruit{"Banana"}, fruit{"Cherry"})
   107  
   108  	pr := ptest.RunAndValidate(t, pipeline)
   109  	checkFruitCount(t, pr)
   110  }
   111  
   112  func TestLPErrorReshufflePipeline_DoFnSideInput(t *testing.T) {
   113  	integration.CheckFilters(t)
   114  
   115  	pipeline, s := beam.NewPipelineWithRoot()
   116  	got := LPErrorReshufflePipeline(s)
   117  	counted := beam.ParDo(s, countFruit, got)
   118  	beam.ParDo0(s, &iterSideStrings{
   119  		Wants: []string{"Apple", "Banana", "Cherry"},
   120  	}, beam.Impulse(s), beam.SideInput{Input: counted})
   121  
   122  	pr := ptest.RunAndValidate(t, pipeline)
   123  	checkFruitCount(t, pr)
   124  }