github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/utils/windowed_value_test.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  """Unit tests for the windowed_value."""
    19  
    20  # pytype: skip-file
    21  
    22  import copy
    23  import itertools
    24  import pickle
    25  import unittest
    26  
    27  from parameterized import parameterized
    28  from parameterized import parameterized_class
    29  
    30  from apache_beam.utils import windowed_value
    31  from apache_beam.utils.timestamp import Timestamp
    32  
    33  
    34  class WindowedValueTest(unittest.TestCase):
    35    def test_timestamps(self):
    36      wv = windowed_value.WindowedValue(None, 3, ())
    37      self.assertEqual(wv.timestamp, Timestamp.of(3))
    38      self.assertTrue(wv.timestamp is wv.timestamp)
    39      self.assertEqual(
    40          windowed_value.WindowedValue(None, -2.5, ()).timestamp,
    41          Timestamp.of(-2.5))
    42  
    43    def test_with_value(self):
    44      pane_info = windowed_value.PaneInfo(
    45          True, True, windowed_value.PaneInfoTiming.ON_TIME, 0, 0)
    46      wv = windowed_value.WindowedValue(1, 3, (), pane_info)
    47      self.assertEqual(
    48          wv.with_value(10), windowed_value.WindowedValue(10, 3, (), pane_info))
    49  
    50    def test_equality(self):
    51      self.assertEqual(
    52          windowed_value.WindowedValue(1, 3, ()),
    53          windowed_value.WindowedValue(1, 3, ()))
    54      self.assertNotEqual(
    55          windowed_value.WindowedValue(1, 3, ()),
    56          windowed_value.WindowedValue(100, 3, ()))
    57      self.assertNotEqual(
    58          windowed_value.WindowedValue(1, 3, ()),
    59          windowed_value.WindowedValue(1, 300, ()))
    60      self.assertNotEqual(
    61          windowed_value.WindowedValue(1, 3, ()),
    62          windowed_value.WindowedValue(1, 300, ((), )))
    63  
    64      self.assertNotEqual(windowed_value.WindowedValue(1, 3, ()), object())
    65  
    66    def test_hash(self):
    67      wv = windowed_value.WindowedValue(1, 3, ())
    68      wv_copy = copy.copy(wv)
    69      self.assertFalse(wv is wv_copy)
    70      self.assertEqual({wv: 100}.get(wv_copy), 100)
    71  
    72    def test_pickle(self):
    73      pane_info = windowed_value.PaneInfo(
    74          True, True, windowed_value.PaneInfoTiming.ON_TIME, 0, 0)
    75      wv = windowed_value.WindowedValue(1, 3, (), pane_info)
    76      self.assertTrue(pickle.loads(pickle.dumps(wv)) == wv)
    77  
    78  
    79  WINDOWED_BATCH_INSTANCES = [
    80      windowed_value.HomogeneousWindowedBatch.of(
    81          None, 3, (), windowed_value.PANE_INFO_UNKNOWN),
    82      windowed_value.HomogeneousWindowedBatch.of(
    83          None,
    84          3, (),
    85          windowed_value.PaneInfo(
    86              True, False, windowed_value.PaneInfoTiming.ON_TIME, 0, 0)),
    87  ]
    88  
    89  
    90  class WindowedBatchTest(unittest.TestCase):
    91    def test_homogeneous_windowed_batch_with_values(self):
    92      pane_info = windowed_value.PaneInfo(
    93          True, True, windowed_value.PaneInfoTiming.ON_TIME, 0, 0)
    94      wb = windowed_value.HomogeneousWindowedBatch.of(['foo', 'bar'],
    95                                                      6, (),
    96                                                      pane_info)
    97      self.assertEqual(
    98          wb.with_values(['baz', 'foo']),
    99          windowed_value.HomogeneousWindowedBatch.of(['baz', 'foo'],
   100                                                     6, (),
   101                                                     pane_info))
   102  
   103    def test_homogeneous_windowed_batch_as_windowed_values(self):
   104      pane_info = windowed_value.PaneInfo(
   105          True, True, windowed_value.PaneInfoTiming.ON_TIME, 0, 0)
   106      wb = windowed_value.HomogeneousWindowedBatch.of(['foo', 'bar'],
   107                                                      3, (),
   108                                                      pane_info)
   109  
   110      self.assertEqual(
   111          list(wb.as_windowed_values(iter)),
   112          [
   113              windowed_value.WindowedValue('foo', 3, (), pane_info),
   114              windowed_value.WindowedValue('bar', 3, (), pane_info)
   115          ])
   116  
   117    @parameterized.expand(itertools.combinations(WINDOWED_BATCH_INSTANCES, 2))
   118    def test_inequality(self, left_wb, right_wb):
   119      self.assertNotEqual(left_wb, right_wb)
   120  
   121    def test_equals_different_type(self):
   122      wb = windowed_value.HomogeneousWindowedBatch.of(
   123          None, 3, (), windowed_value.PANE_INFO_UNKNOWN)
   124      self.assertNotEqual(wb, object())
   125  
   126    def test_homogeneous_from_windowed_values(self):
   127      pane_info = windowed_value.PaneInfo(
   128          True, True, windowed_value.PaneInfoTiming.ON_TIME, 0, 0)
   129  
   130      windowed_values = [
   131          windowed_value.WindowedValue('foofoo', 3, (), pane_info),
   132          windowed_value.WindowedValue('foobar', 6, (), pane_info),
   133          windowed_value.WindowedValue('foobaz', 9, (), pane_info),
   134          windowed_value.WindowedValue('barfoo', 3, (), pane_info),
   135          windowed_value.WindowedValue('barbar', 6, (), pane_info),
   136          windowed_value.WindowedValue('barbaz', 9, (), pane_info),
   137          windowed_value.WindowedValue('bazfoo', 3, (), pane_info),
   138          windowed_value.WindowedValue('bazbar', 6, (), pane_info),
   139          windowed_value.WindowedValue('bazbaz', 9, (), pane_info),
   140      ]
   141  
   142      self.assertEqual(
   143          list(
   144              windowed_value.WindowedBatch.from_windowed_values(
   145                  windowed_values, produce_fn=list)),
   146          [
   147              windowed_value.HomogeneousWindowedBatch.of(
   148                  ['foofoo', 'barfoo', 'bazfoo'], 3, (), pane_info),
   149              windowed_value.HomogeneousWindowedBatch.of(
   150                  ['foobar', 'barbar', 'bazbar'], 6, (), pane_info),
   151              windowed_value.HomogeneousWindowedBatch.of(
   152                  ['foobaz', 'barbaz', 'bazbaz'], 9, (), pane_info)
   153          ])
   154  
   155  
   156  @parameterized_class(('wb', ), [(wb, ) for wb in WINDOWED_BATCH_INSTANCES])
   157  class WindowedBatchUtilitiesTest(unittest.TestCase):
   158    def test_hash(self):
   159      wb_copy = copy.copy(self.wb)
   160      self.assertFalse(self.wb is wb_copy)
   161      self.assertEqual({self.wb: 100}.get(wb_copy), 100)
   162  
   163    def test_pickle(self):
   164      self.assertTrue(pickle.loads(pickle.dumps(self.wb)) == self.wb)
   165  
   166  
   167  if __name__ == '__main__':
   168    unittest.main()