github.com/apache/beam/sdks/v2@v2.48.2/go/test/integration/synthetic/synthetic.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 synthetic contains pipelines for testing synthetic steps and sources.
    17  package synthetic
    18  
    19  import (
    20  	"github.com/apache/beam/sdks/v2/go/pkg/beam"
    21  	"github.com/apache/beam/sdks/v2/go/pkg/beam/io/synthetic"
    22  	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/passert"
    23  )
    24  
    25  // SimplePipeline creates a very simple synthetic pipeline to test that basic
    26  // synthetic pipelines work.
    27  func SimplePipeline() *beam.Pipeline {
    28  	p, s := beam.NewPipelineWithRoot()
    29  	const size = 100
    30  
    31  	src := synthetic.SourceSingle(s,
    32  		synthetic.DefaultSourceConfig().NumElements(size).Build())
    33  	step := synthetic.Step(s, synthetic.DefaultStepConfig().Build(), src)
    34  	passert.Count(s, step, "out", size)
    35  
    36  	return p
    37  }
    38  
    39  // SplittablePipeline creates a simple synthetic pipeline that exercises
    40  // splitting-related behavior.
    41  func SplittablePipeline() *beam.Pipeline {
    42  	p, s := beam.NewPipelineWithRoot()
    43  	const srcSize1 = 50
    44  	const srcSize2 = 10
    45  	const stepMult = 500
    46  	const outCount = (srcSize1 + srcSize2) * stepMult
    47  
    48  	configs := beam.Create(s,
    49  		synthetic.DefaultSourceConfig().NumElements(srcSize1).InitialSplits(3).Build(),
    50  		synthetic.DefaultSourceConfig().NumElements(srcSize2).InitialSplits(3).Build())
    51  	src := synthetic.Source(s, configs)
    52  	step := synthetic.Step(
    53  		s,
    54  		synthetic.
    55  			DefaultStepConfig().
    56  			OutputPerInput(stepMult).
    57  			Splittable(true).
    58  			InitialSplits(8).
    59  			Build(),
    60  		src)
    61  	passert.Count(s, step, "out", outCount)
    62  
    63  	return p
    64  }