github.com/apache/beam/sdks/v2@v2.48.2/go/test/integration/primitives/teststream.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 primitives
    17  
    18  import (
    19  	"github.com/apache/beam/sdks/v2/go/pkg/beam"
    20  	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/passert"
    21  	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/teststream"
    22  )
    23  
    24  // TestStreamSequence tests the TestStream primitive by inserting string elements
    25  // then advancing the watermark past the point where they were inserted.
    26  func TestStreamStrings() *beam.Pipeline {
    27  	p, s := beam.NewPipelineWithRoot()
    28  	con := teststream.NewConfig()
    29  	con.AddElements(100, "a", "b", "c")
    30  	con.AdvanceWatermarkToInfinity()
    31  
    32  	col := teststream.Create(s, con)
    33  
    34  	passert.Count(s, col, "teststream strings", 3)
    35  
    36  	return p
    37  }
    38  
    39  // TestStreamByteSliceSequence tests the TestStream primitive by inserting byte slice elements
    40  // then advancing the watermark to infinity and comparing the output..
    41  func TestStreamByteSliceSequence() *beam.Pipeline {
    42  	p, s := beam.NewPipelineWithRoot()
    43  	con := teststream.NewConfig()
    44  	b := []byte{91, 92, 93}
    45  	con.AddElements(1, b)
    46  	con.AdvanceWatermarkToInfinity()
    47  	col := teststream.Create(s, con)
    48  	passert.Count(s, col, "teststream byte", 1)
    49  	passert.Equals(s, col, append([]byte{3}, b...))
    50  	return p
    51  }
    52  
    53  // TestStreamInt64Sequence tests the TestStream primitive by inserting int64 elements
    54  // then advancing the watermark past the point where they were inserted.
    55  func TestStreamInt64Sequence() *beam.Pipeline {
    56  	p, s := beam.NewPipelineWithRoot()
    57  	con := teststream.NewConfig()
    58  	ele := []int64{91, 92, 93}
    59  	con.AddElementList(100, ele)
    60  	con.AdvanceWatermarkToInfinity()
    61  
    62  	col := teststream.Create(s, con)
    63  
    64  	passert.Count(s, col, "teststream int64", 3)
    65  	passert.EqualsList(s, col, ele)
    66  	return p
    67  }
    68  
    69  // TestStreamTwoInt64Sequences tests the TestStream primitive by inserting two sets of
    70  // int64 elements that arrive on-time into the TestStream
    71  func TestStreamTwoInt64Sequences() *beam.Pipeline {
    72  	p, s := beam.NewPipelineWithRoot()
    73  	con := teststream.NewConfig()
    74  	eo := []int64{91, 92, 93}
    75  	et := []int64{96, 97, 98}
    76  	con.AddElementList(100, eo)
    77  	con.AdvanceWatermark(110)
    78  	con.AddElementList(120, et)
    79  	con.AdvanceWatermark(130)
    80  
    81  	col := teststream.Create(s, con)
    82  
    83  	passert.Count(s, col, "teststream int64", 6)
    84  	passert.EqualsList(s, col, append(eo, et...))
    85  	return p
    86  }
    87  
    88  // TestStreamFloat64Sequence tests the TestStream primitive by inserting float64 elements
    89  // then advancing the watermark past the point where they were inserted.
    90  func TestStreamFloat64Sequence() *beam.Pipeline {
    91  	p, s := beam.NewPipelineWithRoot()
    92  	con := teststream.NewConfig()
    93  	ele := []float64{91.1, 92.2, 93.3}
    94  	con.AddElementList(100, ele)
    95  	con.AdvanceWatermarkToInfinity()
    96  
    97  	col := teststream.Create(s, con)
    98  
    99  	passert.Count(s, col, "teststream float64", 3)
   100  	passert.EqualsList(s, col, ele)
   101  	return p
   102  }
   103  
   104  // TestStreamTwoFloat64Sequences tests the TestStream primitive by inserting two sets of
   105  // float64 elements that arrive on-time into the TestStream
   106  func TestStreamTwoFloat64Sequences() *beam.Pipeline {
   107  	p, s := beam.NewPipelineWithRoot()
   108  	con := teststream.NewConfig()
   109  	eo := []float64{91.1, 92.2, 93.3}
   110  	et := []float64{96.4, 97.5, 98.6}
   111  	con.AddElementList(100, eo)
   112  	con.AdvanceWatermark(110)
   113  	con.AddElementList(120, et)
   114  	con.AdvanceWatermark(130)
   115  
   116  	col := teststream.Create(s, con)
   117  
   118  	passert.Count(s, col, "teststream float64", 6)
   119  	passert.EqualsList(s, col, append(eo, et...))
   120  	return p
   121  }
   122  
   123  // TestStreamBoolSequence tests the TestStream primitive by inserting boolean elements
   124  // then advancing the watermark past the point where they were inserted.
   125  func TestStreamBoolSequence() *beam.Pipeline {
   126  	p, s := beam.NewPipelineWithRoot()
   127  	con := teststream.NewConfig()
   128  	ele := []bool{true, false, true}
   129  	con.AddElementList(100, ele)
   130  	con.AdvanceWatermarkToInfinity()
   131  
   132  	col := teststream.Create(s, con)
   133  
   134  	passert.Count(s, col, "teststream bool", 3)
   135  	passert.EqualsList(s, col, ele)
   136  	return p
   137  }
   138  
   139  // TestStreamTwoBoolSequences tests the TestStream primitive by inserting two sets of
   140  // boolean elements that arrive on-time into the TestStream
   141  func TestStreamTwoBoolSequences() *beam.Pipeline {
   142  	p, s := beam.NewPipelineWithRoot()
   143  	con := teststream.NewConfig()
   144  	eo := []bool{true, false, true}
   145  	et := []bool{false, true, false}
   146  	con.AddElementList(100, eo)
   147  	con.AdvanceWatermark(110)
   148  	con.AddElementList(120, et)
   149  	con.AdvanceWatermark(130)
   150  
   151  	col := teststream.Create(s, con)
   152  
   153  	passert.Count(s, col, "teststream bool", 6)
   154  	passert.EqualsList(s, col, append(eo, et...))
   155  	return p
   156  }