#
# 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.
#

"""Unit tests for our libraries of combine PTransforms."""
# pytype: skip-file

import itertools
import json
import os
import random
import tempfile
import time
import unittest
from pathlib import Path

import hamcrest as hc
import pytest

import apache_beam as beam
import apache_beam.transforms.combiners as combine
from apache_beam import pvalue
from apache_beam.metrics import Metrics
from apache_beam.metrics import MetricsFilter
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import StandardOptions
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.test_stream import TestStream
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
from apache_beam.testing.util import equal_to_per_window
from apache_beam.transforms import WindowInto
from apache_beam.transforms import trigger
from apache_beam.transforms import window
from apache_beam.transforms.core import CombineGlobally
from apache_beam.transforms.core import Create
from apache_beam.transforms.core import Map
from apache_beam.transforms.display import DisplayData
from apache_beam.transforms.display_test import DisplayDataItemMatcher
from apache_beam.transforms.periodicsequence import PeriodicImpulse
from apache_beam.transforms.ptransform import PTransform
from apache_beam.transforms.trigger import AfterAll
from apache_beam.transforms.trigger import AfterCount
from apache_beam.transforms.trigger import AfterWatermark
from apache_beam.transforms.window import FixedWindows
from apache_beam.transforms.window import GlobalWindows
from apache_beam.transforms.window import TimestampCombiner
from apache_beam.transforms.window import TimestampedValue
from apache_beam.typehints import TypeCheckError
from apache_beam.utils.timestamp import Timestamp


class SortedConcatWithCounters(beam.CombineFn):
  """CombineFn for incrementing three different counters:
     counter, distribution, gauge,
     at the same time concatenating words."""
  def __init__(self):
    beam.CombineFn.__init__(self)
    self.word_counter = Metrics.counter(self.__class__, 'word_counter')
    self.word_lengths_counter = Metrics.counter(self.__class__, 'word_lengths')
    self.word_lengths_dist = Metrics.distribution(
        self.__class__, 'word_len_dist')
    self.last_word_len = Metrics.gauge(self.__class__, 'last_word_len')

  def create_accumulator(self):
    return ''

  def add_input(self, acc, element):
    self.word_counter.inc(1)
    self.word_lengths_counter.inc(len(element))
    self.word_lengths_dist.update(len(element))
    self.last_word_len.set(len(element))

    return acc + element

  def merge_accumulators(self, accs):
    return ''.join(accs)

  def extract_output(self, acc):
    # The sorted acc became a list of characters
    # and has to be converted back to a string using join.
    return ''.join(sorted(acc))


class CombineTest(unittest.TestCase):
  def test_builtin_combines(self):
    with TestPipeline() as pipeline:

      vals = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]
      mean = sum(vals) / float(len(vals))
      size = len(vals)
      timestamp = 0

      # First for global combines.
      pcoll = pipeline | 'start' >> Create(vals)
      result_mean = pcoll | 'mean' >> combine.Mean.Globally()
      result_count = pcoll | 'count' >> combine.Count.Globally()
      assert_that(result_mean, equal_to([mean]), label='assert:mean')
      assert_that(result_count, equal_to([size]), label='assert:size')

      # Now for global combines without default
      timestamped = pcoll | Map(lambda x: TimestampedValue(x, timestamp))
      windowed = timestamped | 'window' >> WindowInto(FixedWindows(60))
      result_windowed_mean = (
          windowed
          | 'mean-wo-defaults' >> combine.Mean.Globally().without_defaults())
      assert_that(
          result_windowed_mean,
          equal_to([mean]),
          label='assert:mean-wo-defaults')
      result_windowed_count = (
          windowed
          | 'count-wo-defaults' >> combine.Count.Globally().without_defaults())
      assert_that(
          result_windowed_count,
          equal_to([size]),
          label='assert:count-wo-defaults')

      # Again for per-key combines.
      pcoll = pipeline | 'start-perkey' >> Create([('a', x) for x in vals])
      result_key_mean = pcoll | 'mean-perkey' >> combine.Mean.PerKey()
      result_key_count = pcoll | 'count-perkey' >> combine.Count.PerKey()
      assert_that(result_key_mean, equal_to([('a', mean)]), label='key:mean')
      assert_that(result_key_count, equal_to([('a', size)]), label='key:size')

  def test_top(self):
    with TestPipeline() as pipeline:
      timestamp = 0

      # First for global combines.
      pcoll = pipeline | 'start' >> Create([6, 3, 1, 1, 9, 1, 5, 2, 0, 6])
      result_top = pcoll | 'top' >> combine.Top.Largest(5)
      result_bot = pcoll | 'bot' >> combine.Top.Smallest(4)
      assert_that(result_top, equal_to([[9, 6, 6, 5, 3]]), label='assert:top')
      assert_that(result_bot, equal_to([[0, 1, 1, 1]]), label='assert:bot')

      # Now for global combines without default
      timestamped = pcoll | Map(lambda x: TimestampedValue(x, timestamp))
      windowed = timestamped | 'window' >> WindowInto(FixedWindows(60))
      result_windowed_top = windowed | 'top-wo-defaults' >> combine.Top.Largest(
          5, has_defaults=False)
      result_windowed_bot = (
          windowed
          | 'bot-wo-defaults' >> combine.Top.Smallest(4, has_defaults=False))
      assert_that(
          result_windowed_top,
          equal_to([[9, 6, 6, 5, 3]]),
          label='assert:top-wo-defaults')
      assert_that(
          result_windowed_bot,
          equal_to([[0, 1, 1, 1]]),
          label='assert:bot-wo-defaults')

      # Again for per-key combines.
      pcoll = pipeline | 'start-perkey' >> Create(
          [('a', x) for x in [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]])
      result_key_top = pcoll | 'top-perkey' >> combine.Top.LargestPerKey(5)
      result_key_bot = pcoll | 'bot-perkey' >> combine.Top.SmallestPerKey(4)
      assert_that(
          result_key_top, equal_to([('a', [9, 6, 6, 5, 3])]), label='key:top')
      assert_that(
          result_key_bot, equal_to([('a', [0, 1, 1, 1])]), label='key:bot')

  def test_empty_global_top(self):
    with TestPipeline() as p:
      assert_that(p | beam.Create([]) | combine.Top.Largest(10), equal_to([[]]))

  def test_sharded_top(self):
    elements = list(range(100))
    random.shuffle(elements)

    with TestPipeline() as pipeline:
      shards = [
          pipeline | 'Shard%s' % shard >> beam.Create(elements[shard::7])
          for shard in range(7)
      ]
      assert_that(
          shards | beam.Flatten() | combine.Top.Largest(10),
          equal_to([[99, 98, 97, 96, 95, 94, 93, 92, 91, 90]]))

  def test_top_key(self):
    self.assertEqual(['aa', 'bbb', 'c', 'dddd'] | combine.Top.Of(3, key=len),
                     [['dddd', 'bbb', 'aa']])
    self.assertEqual(['aa', 'bbb', 'c', 'dddd']
                     | combine.Top.Of(3, key=len, reverse=True),
                     [['c', 'aa', 'bbb']])

    self.assertEqual(['xc', 'zb', 'yd', 'wa']
                     | combine.Top.Largest(3, key=lambda x: x[-1]),
                     [['yd', 'xc', 'zb']])
    self.assertEqual(['xc', 'zb', 'yd', 'wa']
                     | combine.Top.Smallest(3, key=lambda x: x[-1]),
                     [['wa', 'zb', 'xc']])

    self.assertEqual([('a', x) for x in [1, 2, 3, 4, 1, 1]]
                     | combine.Top.LargestPerKey(3, key=lambda x: -x),
                     [('a', [1, 1, 1])])
    self.assertEqual([('a', x) for x in [1, 2, 3, 4, 1, 1]]
                     | combine.Top.SmallestPerKey(3, key=lambda x: -x),
                     [('a', [4, 3, 2])])

  def test_sharded_top_combine_fn(self):
    def test_combine_fn(combine_fn, shards, expected):
      accumulators = [
          combine_fn.add_inputs(combine_fn.create_accumulator(), shard)
          for shard in shards
      ]
      final_accumulator = combine_fn.merge_accumulators(accumulators)
      self.assertEqual(combine_fn.extract_output(final_accumulator), expected)

    test_combine_fn(combine.TopCombineFn(3), [range(10), range(10)], [9, 9, 8])
    test_combine_fn(
        combine.TopCombineFn(5), [range(1000), range(100), range(1001)],
        [1000, 999, 999, 998, 998])

  def test_combine_per_key_top_display_data(self):
    def individual_test_per_key_dd(combineFn):
      transform = beam.CombinePerKey(combineFn)
      dd = DisplayData.create_from(transform)
      expected_items = [
          DisplayDataItemMatcher('combine_fn', combineFn.__class__),
          DisplayDataItemMatcher('n', combineFn._n),
          DisplayDataItemMatcher('compare', combineFn._compare.__name__)
      ]
      hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))

    individual_test_per_key_dd(combine.Largest(5))
    individual_test_per_key_dd(combine.Smallest(3))
    individual_test_per_key_dd(combine.TopCombineFn(8))
    individual_test_per_key_dd(combine.Largest(5))

  def test_combine_sample_display_data(self):
    def individual_test_per_key_dd(sampleFn, n):
      trs = [sampleFn(n)]
      for transform in trs:
        dd = DisplayData.create_from(transform)
        hc.assert_that(
            dd.items,
            hc.contains_inanyorder(DisplayDataItemMatcher('n', transform._n)))

    individual_test_per_key_dd(combine.Sample.FixedSizePerKey, 5)
    individual_test_per_key_dd(combine.Sample.FixedSizeGlobally, 5)

  def test_combine_globally_display_data(self):
    transform = beam.CombineGlobally(combine.Smallest(5))
    dd = DisplayData.create_from(transform)
    expected_items = [
        DisplayDataItemMatcher('combine_fn', combine.Smallest),
        DisplayDataItemMatcher('n', 5),
        DisplayDataItemMatcher('compare', 'gt')
    ]
    hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))

  def test_basic_combiners_display_data(self):
    transform = beam.CombineGlobally(
        combine.TupleCombineFn(max, combine.MeanCombineFn(), sum))
    dd = DisplayData.create_from(transform)
    expected_items = [
        DisplayDataItemMatcher('combine_fn', combine.TupleCombineFn),
        DisplayDataItemMatcher('combiners', "['max', 'MeanCombineFn', 'sum']"),
        DisplayDataItemMatcher('merge_accumulators_batch_size', 333),
    ]
    hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))

  def test_top_shorthands(self):
    with TestPipeline() as pipeline:

      pcoll = pipeline | 'start' >> Create([6, 3, 1, 1, 9, 1, 5, 2, 0, 6])
      result_top = pcoll | 'top' >> beam.CombineGlobally(combine.Largest(5))
      result_bot = pcoll | 'bot' >> beam.CombineGlobally(combine.Smallest(4))
      assert_that(result_top, equal_to([[9, 6, 6, 5, 3]]), label='assert:top')
      assert_that(result_bot, equal_to([[0, 1, 1, 1]]), label='assert:bot')

      pcoll = pipeline | 'start-perkey' >> Create(
          [('a', x) for x in [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]])
      result_ktop = pcoll | 'top-perkey' >> beam.CombinePerKey(
          combine.Largest(5))
      result_kbot = pcoll | 'bot-perkey' >> beam.CombinePerKey(
          combine.Smallest(4))
      assert_that(result_ktop, equal_to([('a', [9, 6, 6, 5, 3])]), label='ktop')
      assert_that(result_kbot, equal_to([('a', [0, 1, 1, 1])]), label='kbot')

  def test_top_no_compact(self):
    class TopCombineFnNoCompact(combine.TopCombineFn):
      def compact(self, accumulator):
        return accumulator

    with TestPipeline() as pipeline:
      pcoll = pipeline | 'Start' >> Create([6, 3, 1, 1, 9, 1, 5, 2, 0, 6])
      result_top = pcoll | 'Top' >> beam.CombineGlobally(
          TopCombineFnNoCompact(5, key=lambda x: x))
      result_bot = pcoll | 'Bot' >> beam.CombineGlobally(
          TopCombineFnNoCompact(4, reverse=True))
      assert_that(result_top, equal_to([[9, 6, 6, 5, 3]]), label='Assert:Top')
      assert_that(result_bot, equal_to([[0, 1, 1, 1]]), label='Assert:Bot')

      pcoll = pipeline | 'Start-Perkey' >> Create(
          [('a', x) for x in [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]])
      result_ktop = pcoll | 'Top-PerKey' >> beam.CombinePerKey(
          TopCombineFnNoCompact(5, key=lambda x: x))
      result_kbot = pcoll | 'Bot-PerKey' >> beam.CombinePerKey(
          TopCombineFnNoCompact(4, reverse=True))
      assert_that(result_ktop, equal_to([('a', [9, 6, 6, 5, 3])]), label='KTop')
      assert_that(result_kbot, equal_to([('a', [0, 1, 1, 1])]), label='KBot')

  def test_global_sample(self):
    def is_good_sample(actual):
      assert len(actual) == 1
      assert sorted(actual[0]) in [[1, 1, 2], [1, 2, 2]], actual

    with TestPipeline() as pipeline:
      timestamp = 0
      pcoll = pipeline | 'start' >> Create([1, 1, 2, 2])

      # Now for global combines without default
      timestamped = pcoll | Map(lambda x: TimestampedValue(x, timestamp))
      windowed = timestamped | 'window' >> WindowInto(FixedWindows(60))

      for ix in range(9):
        assert_that(
            pcoll | 'sample-%d' % ix >> combine.Sample.FixedSizeGlobally(3),
            is_good_sample,
            label='check-%d' % ix)
        result_windowed = (
            windowed
            | 'sample-wo-defaults-%d' % ix >>
            combine.Sample.FixedSizeGlobally(3).without_defaults())
        assert_that(
            result_windowed, is_good_sample, label='check-wo-defaults-%d' % ix)

  def test_per_key_sample(self):
    with TestPipeline() as pipeline:
      pcoll = pipeline | 'start-perkey' >> Create(
          sum(([(i, 1), (i, 1), (i, 2), (i, 2)] for i in range(9)), []))
      result = pcoll | 'sample' >> combine.Sample.FixedSizePerKey(3)

      def matcher():
        def match(actual):
          for _, samples in actual:
            equal_to([3])([len(samples)])
            num_ones = sum(1 for x in samples if x == 1)
            num_twos = sum(1 for x in samples if x == 2)
            equal_to([1, 2])([num_ones, num_twos])

        return match

      assert_that(result, matcher())

  def test_tuple_combine_fn(self):
    with TestPipeline() as p:
      result = (
          p
          | Create([('a', 100, 0.0), ('b', 10, -1), ('c', 1, 100)])
          | beam.CombineGlobally(
              combine.TupleCombineFn(max, combine.MeanCombineFn(),
                                     sum)).without_defaults())
      assert_that(result, equal_to([('c', 111.0 / 3, 99.0)]))

  def test_tuple_combine_fn_without_defaults(self):
    with TestPipeline() as p:
      result = (
          p
          | Create([1, 1, 2, 3])
          | beam.CombineGlobally(
              combine.TupleCombineFn(
                  min, combine.MeanCombineFn(),
                  max).with_common_input()).without_defaults())
      assert_that(result, equal_to([(1, 7.0 / 4, 3)]))

  def test_empty_tuple_combine_fn(self):
    with TestPipeline() as p:
      result = (
          p
          | Create([(), (), ()])
          | beam.CombineGlobally(combine.TupleCombineFn()))
      assert_that(result, equal_to([()]))

  def test_tuple_combine_fn_batched_merge(self):
    num_combine_fns = 10
    max_num_accumulators_in_memory = 30
    # Maximum number of accumulator tuples in memory - 1 for the merge result.
    merge_accumulators_batch_size = (
        max_num_accumulators_in_memory // num_combine_fns - 1)
    num_accumulator_tuples_to_merge = 20

    class CountedAccumulator:
      count = 0
      oom = False

      def __init__(self):
        if CountedAccumulator.count > max_num_accumulators_in_memory:
          CountedAccumulator.oom = True
        else:
          CountedAccumulator.count += 1

    class CountedAccumulatorCombineFn(beam.CombineFn):
      def create_accumulator(self):
        return CountedAccumulator()

      def merge_accumulators(self, accumulators):
        CountedAccumulator.count += 1
        for _ in accumulators:
          CountedAccumulator.count -= 1

    combine_fn = combine.TupleCombineFn(
        *[CountedAccumulatorCombineFn() for _ in range(num_combine_fns)],
        merge_accumulators_batch_size=merge_accumulators_batch_size)
    combine_fn.merge_accumulators(
        combine_fn.create_accumulator()
        for _ in range(num_accumulator_tuples_to_merge))
    assert not CountedAccumulator.oom

  def test_to_list_and_to_dict1(self):
    with TestPipeline() as pipeline:
      the_list = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]
      timestamp = 0
      pcoll = pipeline | 'start' >> Create(the_list)
      result = pcoll | 'to list' >> combine.ToList()

      # Now for global combines without default
      timestamped = pcoll | Map(lambda x: TimestampedValue(x, timestamp))
      windowed = timestamped | 'window' >> WindowInto(FixedWindows(60))
      result_windowed = (
          windowed
          | 'to list wo defaults' >> combine.ToList().without_defaults())

      def matcher(expected):
        def match(actual):
          equal_to(expected[0])(actual[0])

        return match

      assert_that(result, matcher([the_list]))
      assert_that(
          result_windowed, matcher([the_list]), label='to-list-wo-defaults')

  def test_to_list_and_to_dict2(self):
    with TestPipeline() as pipeline:
      pairs = [(1, 2), (3, 4), (5, 6)]
      timestamp = 0
      pcoll = pipeline | 'start-pairs' >> Create(pairs)
      result = pcoll | 'to dict' >> combine.ToDict()

      # Now for global combines without default
      timestamped = pcoll | Map(lambda x: TimestampedValue(x, timestamp))
      windowed = timestamped | 'window' >> WindowInto(FixedWindows(60))
      result_windowed = (
          windowed
          | 'to dict wo defaults' >> combine.ToDict().without_defaults())

      def matcher():
        def match(actual):
          equal_to([1])([len(actual)])
          equal_to(pairs)(actual[0].items())

        return match

      assert_that(result, matcher())
      assert_that(result_windowed, matcher(), label='to-dict-wo-defaults')

  def test_to_set(self):
    pipeline = TestPipeline()
    the_list = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]
    timestamp = 0
    pcoll = pipeline | 'start' >> Create(the_list)
    result = pcoll | 'to set' >> combine.ToSet()

    # Now for global combines without default
    timestamped = pcoll | Map(lambda x: TimestampedValue(x, timestamp))
    windowed = timestamped | 'window' >> WindowInto(FixedWindows(60))
    result_windowed = (
        windowed
        | 'to set wo defaults' >> combine.ToSet().without_defaults())

    def matcher(expected):
      def match(actual):
        equal_to(expected[0])(actual[0])

      return match

    assert_that(result, matcher(set(the_list)))
    assert_that(
        result_windowed, matcher(set(the_list)), label='to-set-wo-defaults')

  def test_combine_globally_with_default(self):
    with TestPipeline() as p:
      assert_that(p | Create([]) | CombineGlobally(sum), equal_to([0]))

  def test_combine_globally_without_default(self):
    with TestPipeline() as p:
      result = p | Create([]) | CombineGlobally(sum).without_defaults()
      assert_that(result, equal_to([]))

  def test_combine_globally_with_default_side_input(self):
    class SideInputCombine(PTransform):
      def expand(self, pcoll):
        side = pcoll | CombineGlobally(sum).as_singleton_view()
        main = pcoll.pipeline | Create([None])
        return main | Map(lambda _, s: s, side)

    with TestPipeline() as p:
      result1 = p | 'i1' >> Create([]) | 'c1' >> SideInputCombine()
      result2 = p | 'i2' >> Create([1, 2, 3, 4]) | 'c2' >> SideInputCombine()
      assert_that(result1, equal_to([0]), label='r1')
      assert_that(result2, equal_to([10]), label='r2')

  def test_hot_key_fanout(self):
    with TestPipeline() as p:
      result = (
          p
          | beam.Create(itertools.product(['hot', 'cold'], range(10)))
          | beam.CombinePerKey(combine.MeanCombineFn()).with_hot_key_fanout(
              lambda key: (key == 'hot') * 5))
      assert_that(result, equal_to([('hot', 4.5), ('cold', 4.5)]))

  def test_hot_key_fanout_sharded(self):
    # Lots of elements with the same key with varying/no fanout.
    with TestPipeline() as p:
      elements = [(None, e) for e in range(1000)]
      random.shuffle(elements)
      shards = [
          p | "Shard%s" % shard >> beam.Create(elements[shard::20])
          for shard in range(20)
      ]
      result = (
          shards
          | beam.Flatten()
          | beam.CombinePerKey(combine.MeanCombineFn()).with_hot_key_fanout(
              lambda key: random.randrange(0, 5)))
      assert_that(result, equal_to([(None, 499.5)]))

  def test_global_fanout(self):
    with TestPipeline() as p:
      result = (
          p
          | beam.Create(range(100))
          | beam.CombineGlobally(combine.MeanCombineFn()).with_fanout(11))
      assert_that(result, equal_to([49.5]))

  def test_combining_with_accumulation_mode_and_fanout(self):
    # PCollection will contain elements from 1 to 5.
    elements = [i for i in range(1, 6)]

    ts = TestStream().advance_watermark_to(0)
    for i in elements:
      ts.add_elements([i])
    ts.advance_watermark_to_infinity()

    options = PipelineOptions()
    options.view_as(StandardOptions).streaming = True
    with TestPipeline(options=options) as p:
      result = (
          p
          | ts
          | beam.WindowInto(
              GlobalWindows(),
              accumulation_mode=trigger.AccumulationMode.ACCUMULATING,
              trigger=AfterWatermark(early=AfterAll(AfterCount(1))))
          | beam.CombineGlobally(sum).without_defaults().with_fanout(2))

      def has_expected_values(actual):
        from hamcrest.core import assert_that as hamcrest_assert
        from hamcrest.library.collection import contains
        from hamcrest.library.collection import only_contains
        ordered = sorted(actual)
        # Early firings.
        hamcrest_assert(ordered[:4], contains(1, 3, 6, 10))
        # Different runners have different number of 15s, but there should
        # be at least one 15.
        hamcrest_assert(ordered[4:], only_contains(15))

      assert_that(result, has_expected_values)

  def test_combining_with_sliding_windows_and_fanout_raises_error(self):
    options = PipelineOptions()
    options.view_as(StandardOptions).streaming = True
    with self.assertRaises(ValueError):
      with TestPipeline(options=options) as p:
        _ = (
            p
            | beam.Create([
                window.TimestampedValue(0, Timestamp(seconds=1666707510)),
                window.TimestampedValue(1, Timestamp(seconds=1666707511)),
                window.TimestampedValue(2, Timestamp(seconds=1666707512)),
                window.TimestampedValue(3, Timestamp(seconds=1666707513)),
                window.TimestampedValue(5, Timestamp(seconds=1666707515)),
                window.TimestampedValue(6, Timestamp(seconds=1666707516)),
                window.TimestampedValue(7, Timestamp(seconds=1666707517)),
                window.TimestampedValue(8, Timestamp(seconds=1666707518))
            ])
            | beam.WindowInto(window.SlidingWindows(10, 5))
            | beam.CombineGlobally(beam.combiners.ToListCombineFn()).
            without_defaults().with_fanout(7))

  def test_MeanCombineFn_combine(self):
    with TestPipeline() as p:
      input = (
          p
          | beam.Create([('a', 1), ('a', 1), ('a', 4), ('b', 1), ('b', 13)]))
      # The mean of all values regardless of key.
      global_mean = (
          input
          | beam.Values()
          | beam.CombineGlobally(combine.MeanCombineFn()))

      # The (key, mean) pairs for all keys.
      mean_per_key = (input | beam.CombinePerKey(combine.MeanCombineFn()))

      expected_mean_per_key = [('a', 2), ('b', 7)]
      assert_that(global_mean, equal_to([4]), label='global mean')
      assert_that(
          mean_per_key, equal_to(expected_mean_per_key), label='mean per key')

  def test_MeanCombineFn_combine_empty(self):
    # For each element in a PCollection, if it is float('NaN'), then emits
    # a string 'NaN', otherwise emits str(element).

    with TestPipeline() as p:
      input = (p | beam.Create([]))

      # Compute the mean of all values in the PCollection,
      # then format the mean. Since the Pcollection is empty,
      # the mean is float('NaN'), and is formatted to be a string 'NaN'.
      global_mean = (
          input
          | beam.Values()
          | beam.CombineGlobally(combine.MeanCombineFn())
          | beam.Map(str))

      mean_per_key = (input | beam.CombinePerKey(combine.MeanCombineFn()))

      # We can't compare one float('NaN') with another float('NaN'),
      # but we can compare one 'nan' string with another string.
      assert_that(global_mean, equal_to(['nan']), label='global mean')
      assert_that(mean_per_key, equal_to([]), label='mean per key')

  def test_sessions_combine(self):
    with TestPipeline() as p:
      input = (
          p
          | beam.Create([('c', 1), ('c', 9), ('c', 12), ('d', 2), ('d', 4)])
          | beam.MapTuple(lambda k, v: window.TimestampedValue((k, v), v))
          | beam.WindowInto(window.Sessions(4)))

      global_sum = (
          input
          | beam.Values()
          | beam.CombineGlobally(sum).without_defaults())
      sum_per_key = input | beam.CombinePerKey(sum)

      # The first window has 3 elements: ('c', 1), ('d', 2), ('d', 4).
      # The second window has 2 elements: ('c', 9), ('c', 12).
      assert_that(global_sum, equal_to([7, 21]), label='global sum')
      assert_that(
          sum_per_key,
          equal_to([('c', 1), ('c', 21), ('d', 6)]),
          label='sum per key')

  def test_fixed_windows_combine(self):
    with TestPipeline() as p:
      input = (
          p
          | beam.Create([('c', 1), ('c', 2), ('c', 10), ('d', 5), ('d', 8),
                         ('d', 9)])
          | beam.MapTuple(lambda k, v: window.TimestampedValue((k, v), v))
          | beam.WindowInto(window.FixedWindows(4)))

      global_sum = (
          input
          | beam.Values()
          | beam.CombineGlobally(sum).without_defaults())
      sum_per_key = input | beam.CombinePerKey(sum)

      # The first window has 2 elements: ('c', 1), ('c', 2).
      # The second window has 1 elements: ('d', 5).
      # The third window has 3 elements: ('c', 10), ('d', 8), ('d', 9).
      assert_that(global_sum, equal_to([3, 5, 27]), label='global sum')
      assert_that(
          sum_per_key,
          equal_to([('c', 3), ('c', 10), ('d', 5), ('d', 17)]),
          label='sum per key')

  # Test that three different kinds of metrics work with a customized
  # SortedConcatWithCounters CombineFn.
  def test_custormized_counters_in_combine_fn(self):
    p = TestPipeline()
    input = (
        p
        | beam.Create([('key1', 'a'), ('key1', 'ab'), ('key1', 'abc'),
                       ('key2', 'uvxy'), ('key2', 'uvxyz')]))

    # The result of concatenating all values regardless of key.
    global_concat = (
        input
        | beam.Values()
        | beam.CombineGlobally(SortedConcatWithCounters()))

    # The (key, concatenated_string) pairs for all keys.
    concat_per_key = (input | beam.CombinePerKey(SortedConcatWithCounters()))

    # Verify the concatenated strings are correct.
    expected_concat_per_key = [('key1', 'aaabbc'), ('key2', 'uuvvxxyyz')]
    assert_that(
        global_concat, equal_to(['aaabbcuuvvxxyyz']), label='global concat')
    assert_that(
        concat_per_key,
        equal_to(expected_concat_per_key),
        label='concat per key')

    result = p.run()
    result.wait_until_finish()

    # Verify the values of metrics are correct.
    word_counter_filter = MetricsFilter().with_name('word_counter')
    query_result = result.metrics().query(word_counter_filter)
    if query_result['counters']:
      word_counter = query_result['counters'][0]
      self.assertEqual(word_counter.result, 5)

    word_lengths_filter = MetricsFilter().with_name('word_lengths')
    query_result = result.metrics().query(word_lengths_filter)
    if query_result['counters']:
      word_lengths = query_result['counters'][0]
      self.assertEqual(word_lengths.result, 15)

    word_len_dist_filter = MetricsFilter().with_name('word_len_dist')
    query_result = result.metrics().query(word_len_dist_filter)
    if query_result['distributions']:
      word_len_dist = query_result['distributions'][0]
      self.assertEqual(word_len_dist.result.mean, 3)

    last_word_len_filter = MetricsFilter().with_name('last_word_len')
    query_result = result.metrics().query(last_word_len_filter)
    if query_result['gauges']:
      last_word_len = query_result['gauges'][0]
      self.assertIn(last_word_len.result.value, [1, 2, 3, 4, 5])

  # Test that three different kinds of metrics work with the customized
  # SortedConcatWithCounters CombineFn when the PCollection is empty.
  def test_custormized_counters_in_combine_fn_empty(self):
    p = TestPipeline()
    input = p | beam.Create([])

    # The result of concatenating all values regardless of key.
    global_concat = (
        input
        | beam.Values()
        | beam.CombineGlobally(SortedConcatWithCounters()))

    # The (key, concatenated_string) pairs for all keys.
    concat_per_key = (input | beam.CombinePerKey(SortedConcatWithCounters()))

    # Verify the concatenated strings are correct.
    assert_that(global_concat, equal_to(['']), label='global concat')
    assert_that(concat_per_key, equal_to([]), label='concat per key')

    result = p.run()
    result.wait_until_finish()

    # Verify the values of metrics are correct.
    word_counter_filter = MetricsFilter().with_name('word_counter')
    query_result = result.metrics().query(word_counter_filter)
    if query_result['counters']:
      word_counter = query_result['counters'][0]
      self.assertEqual(word_counter.result, 0)

    word_lengths_filter = MetricsFilter().with_name('word_lengths')
    query_result = result.metrics().query(word_lengths_filter)
    if query_result['counters']:
      word_lengths = query_result['counters'][0]
      self.assertEqual(word_lengths.result, 0)

    word_len_dist_filter = MetricsFilter().with_name('word_len_dist')
    query_result = result.metrics().query(word_len_dist_filter)
    if query_result['distributions']:
      word_len_dist = query_result['distributions'][0]
      self.assertEqual(word_len_dist.result.count, 0)

    last_word_len_filter = MetricsFilter().with_name('last_word_len')
    query_result = result.metrics().query(last_word_len_filter)

    # No element has ever been recorded.
    self.assertFalse(query_result['gauges'])


class LatestTest(unittest.TestCase):
  def test_globally(self):
    l = [
        window.TimestampedValue(3, 100),
        window.TimestampedValue(1, 200),
        window.TimestampedValue(2, 300)
    ]
    with TestPipeline() as p:
      # Map(lambda x: x) PTransform is added after Create here, because when
      # a PCollection of TimestampedValues is created with Create PTransform,
      # the timestamps are not assigned to it. Adding a Map forces the
      # PCollection to go through a DoFn so that the PCollection consists of
      # the elements with timestamps assigned to them instead of a PCollection
      # of TimestampedValue(element, timestamp).
      pcoll = p | Create(l) | Map(lambda x: x)
      latest = pcoll | combine.Latest.Globally()
      assert_that(latest, equal_to([2]))

      # Now for global combines without default
      windowed = pcoll | 'window' >> WindowInto(FixedWindows(180))
      result_windowed = (
          windowed
          |
          'latest wo defaults' >> combine.Latest.Globally().without_defaults())

      assert_that(result_windowed, equal_to([3, 2]), label='latest-wo-defaults')

  def test_globally_empty(self):
    l = []
    with TestPipeline() as p:
      pc = p | Create(l) | Map(lambda x: x)
      latest = pc | combine.Latest.Globally()
      assert_that(latest, equal_to([None]))

  def test_per_key(self):
    l = [
        window.TimestampedValue(('a', 1), 300),
        window.TimestampedValue(('b', 3), 100),
        window.TimestampedValue(('a', 2), 200)
    ]
    with TestPipeline() as p:
      pc = p | Create(l) | Map(lambda x: x)
      latest = pc | combine.Latest.PerKey()
      assert_that(latest, equal_to([('a', 1), ('b', 3)]))

  def test_per_key_empty(self):
    l = []
    with TestPipeline() as p:
      pc = p | Create(l) | Map(lambda x: x)
      latest = pc | combine.Latest.PerKey()
      assert_that(latest, equal_to([]))


class LatestCombineFnTest(unittest.TestCase):
  def setUp(self):
    self.fn = combine.LatestCombineFn()

  def test_create_accumulator(self):
    accumulator = self.fn.create_accumulator()
    self.assertEqual(accumulator, (None, window.MIN_TIMESTAMP))

  def test_add_input(self):
    accumulator = self.fn.create_accumulator()
    element = (1, 100)
    new_accumulator = self.fn.add_input(accumulator, element)
    self.assertEqual(new_accumulator, (1, 100))

  def test_merge_accumulators(self):
    accumulators = [(2, 400), (5, 100), (9, 200)]
    merged_accumulator = self.fn.merge_accumulators(accumulators)
    self.assertEqual(merged_accumulator, (2, 400))

  def test_extract_output(self):
    accumulator = (1, 100)
    output = self.fn.extract_output(accumulator)
    self.assertEqual(output, 1)

  def test_with_input_types_decorator_violation(self):
    l_int = [1, 2, 3]
    l_dict = [{'a': 3}, {'g': 5}, {'r': 8}]
    l_3_tuple = [(12, 31, 41), (12, 34, 34), (84, 92, 74)]

    with self.assertRaises(TypeCheckError):
      with TestPipeline() as p:
        pc = p | Create(l_int)
        _ = pc | beam.CombineGlobally(self.fn)

    with self.assertRaises(TypeCheckError):
      with TestPipeline() as p:
        pc = p | Create(l_dict)
        _ = pc | beam.CombineGlobally(self.fn)

    with self.assertRaises(TypeCheckError):
      with TestPipeline() as p:
        pc = p | Create(l_3_tuple)
        _ = pc | beam.CombineGlobally(self.fn)


@pytest.mark.it_validatesrunner
class CombineValuesTest(unittest.TestCase):
  def test_gbk_immediately_followed_by_combine(self):
    def merge(vals):
      return "".join(vals)

    with TestPipeline() as p:
      result = (
          p \
          | Create([("key1", "foo"), ("key2", "bar"), ("key1", "foo")],
                    reshuffle=False) \
          | beam.GroupByKey() \
          | beam.CombineValues(merge) \
          | beam.MapTuple(lambda k, v: '{}: {}'.format(k, v)))

      assert_that(result, equal_to(['key1: foofoo', 'key2: bar']))


#
# Test cases for streaming.
#
@pytest.mark.it_validatesrunner
class TimestampCombinerTest(unittest.TestCase):
  def test_combiner_earliest(self):
    """Test TimestampCombiner with EARLIEST."""
    options = PipelineOptions(streaming=True)
    with TestPipeline(options=options) as p:
      result = (
          p
          | TestStream().add_elements([window.TimestampedValue(
              ('k', 100), 2)]).add_elements(
                  [window.TimestampedValue(
                      ('k', 400), 7)]).advance_watermark_to_infinity()
          | beam.WindowInto(
              window.FixedWindows(10),
              timestamp_combiner=TimestampCombiner.OUTPUT_AT_EARLIEST)
          | beam.CombinePerKey(sum))

      records = (
          result
          | beam.Map(lambda e, ts=beam.DoFn.TimestampParam: (e, ts)))

      # All the KV pairs are applied GBK using EARLIEST timestamp for the same
      # key.
      expected_window_to_elements = {
          window.IntervalWindow(0, 10): [
              (('k', 500), Timestamp(2)),
          ],
      }

      assert_that(
          records,
          equal_to_per_window(expected_window_to_elements),
          use_global_window=False,
          label='assert per window')

  def test_combiner_latest(self):
    """Test TimestampCombiner with LATEST."""
    options = PipelineOptions(streaming=True)
    with TestPipeline(options=options) as p:
      result = (
          p
          | TestStream().add_elements([window.TimestampedValue(
              ('k', 100), 2)]).add_elements(
                  [window.TimestampedValue(
                      ('k', 400), 7)]).advance_watermark_to_infinity()
          | beam.WindowInto(
              window.FixedWindows(10),
              timestamp_combiner=TimestampCombiner.OUTPUT_AT_LATEST)
          | beam.CombinePerKey(sum))

      records = (
          result
          | beam.Map(lambda e, ts=beam.DoFn.TimestampParam: (e, ts)))

      # All the KV pairs are applied GBK using LATEST timestamp for
      # the same key.
      expected_window_to_elements = {
          window.IntervalWindow(0, 10): [
              (('k', 500), Timestamp(7)),
          ],
      }

      assert_that(
          records,
          equal_to_per_window(expected_window_to_elements),
          use_global_window=False,
          label='assert per window')


class CombineGloballyTest(unittest.TestCase):
  def test_combine_globally_for_unbounded_source_with_default(self):
    # this error is logged since the below combination is ill-defined.
    with self.assertLogs() as captured_logs:
      with TestPipeline() as p:
        _ = (
            p
            | PeriodicImpulse(
                start_timestamp=time.time(),
                stop_timestamp=time.time() + 4,
                fire_interval=1,
                apply_windowing=False,
            )
            | beam.Map(lambda x: ('c', 1))
            | beam.WindowInto(
                window.GlobalWindows(),
                trigger=trigger.Repeatedly(trigger.AfterCount(2)),
                accumulation_mode=trigger.AccumulationMode.DISCARDING,
            )
            | beam.combiners.Count.Globally())
    self.assertIn('unbounded collections', '\n'.join(captured_logs.output))

  def test_combine_globally_for_unbounded_source_without_defaults(self):
    # this is the supported case
    with TestPipeline() as p:
      _ = (
          p
          | PeriodicImpulse(
              start_timestamp=time.time(),
              stop_timestamp=time.time() + 4,
              fire_interval=1,
              apply_windowing=False,
          )
          | beam.Map(lambda x: 1)
          | beam.WindowInto(
              window.GlobalWindows(),
              trigger=trigger.Repeatedly(trigger.AfterCount(2)),
              accumulation_mode=trigger.AccumulationMode.DISCARDING,
          )
          | beam.CombineGlobally(sum).without_defaults())


def get_common_items(sets, excluded_chars=""):
  # set.intersection() takes multiple sets as separete arguments.
  # We unpack the `sets` list into multiple arguments with the * operator.
  # The combine transform might give us an empty list of `sets`,
  # so we use a list with an empty set as a default value.
  common = set.intersection(*(sets or [set()]))
  return common.difference(excluded_chars)


class CombinerWithSideInputs(unittest.TestCase):
  def test_cpk_with_side_input(self):
    test_cases = [(get_common_items, True),
                  (beam.CombineFn.from_callable(get_common_items), True),
                  (get_common_items, False),
                  (beam.CombineFn.from_callable(get_common_items), False)]
    for combiner, with_kwarg in test_cases:
      self._check_combineperkey_with_side_input(combiner, with_kwarg)
      self._check_combineglobally_with_side_input(combiner, with_kwarg)

  def _check_combineperkey_with_side_input(self, combiner, with_kwarg):
    with beam.Pipeline() as pipeline:
      pc = (pipeline | beam.Create(['🍅']))
      if with_kwarg:
        cpk = beam.CombinePerKey(
            combiner, excluded_chars=beam.pvalue.AsSingleton(pc))
      else:
        cpk = beam.CombinePerKey(combiner, beam.pvalue.AsSingleton(pc))
      common_items = (
          pipeline
          | 'Create produce' >> beam.Create([
              {'🍓', '🥕', '🍌', '🍅', '🌶️'},
              {'🍇', '🥕', '🥝', '🍅', '🥔'},
              {'🍉', '🥕', '🍆', '🍅', '🍍'},
              {'🥑', '🥕', '🌽', '🍅', '🥥'},
          ])
          | beam.WithKeys(lambda x: None)
          | cpk)
      assert_that(common_items, equal_to([(None, {'🥕'})]))

  def _check_combineglobally_with_side_input(self, combiner, with_kwarg):
    with beam.Pipeline() as pipeline:
      pc = (pipeline | beam.Create(['🍅']))
      if with_kwarg:
        cpk = beam.CombineGlobally(
            combiner, excluded_chars=beam.pvalue.AsSingleton(pc))
      else:
        cpk = beam.CombineGlobally(combiner, beam.pvalue.AsSingleton(pc))
      common_items = (
          pipeline
          | 'Create produce' >> beam.Create([
              {'🍓', '🥕', '🍌', '🍅', '🌶️'},
              {'🍇', '🥕', '🥝', '🍅', '🥔'},
              {'🍉', '🥕', '🍆', '🍅', '🍍'},
              {'🥑', '🥕', '🌽', '🍅', '🥥'},
          ])
          | cpk)
      assert_that(common_items, equal_to([{'🥕'}]))

  def test_combinefn_methods_with_side_input(self):
    # Test that the expected combinefn methods are called with the
    # expected arguments when using side inputs in CombinePerKey.
    with tempfile.TemporaryDirectory() as tmp_dirname:
      fname = str(Path(tmp_dirname) / "combinefn_calls.json")
      with open(fname, "w") as f:
        json.dump({}, f)

      def set_in_json(key, values):
        current_json = {}
        if os.path.exists(fname):
          with open(fname, "r") as f:
            current_json = json.load(f)
        current_json[key] = values
        with open(fname, "w") as f:
          json.dump(current_json, f)

      class MyCombiner(beam.CombineFn):
        def create_accumulator(self, *args, **kwargs):
          set_in_json("create_accumulator_args", args)
          set_in_json("create_accumulator_kwargs", kwargs)
          return args, kwargs

        def add_input(self, accumulator, input, *args, **kwargs):
          set_in_json("add_input_args", args)
          set_in_json("add_input_kwargs", kwargs)
          return accumulator

        def merge_accumulators(self, accumulators, *args, **kwargs):
          set_in_json("merge_accumulators_args", args)
          set_in_json("merge_accumulators_kwargs", kwargs)
          return args, kwargs

        def compact(self, accumulator, *args, **kwargs):
          set_in_json("compact_args", args)
          set_in_json("compact_kwargs", kwargs)
          return accumulator

        def extract_output(self, accumulator, *args, **kwargs):
          set_in_json("extract_output_args", args)
          set_in_json("extract_output_kwargs", kwargs)
          return accumulator

      with beam.Pipeline() as p:
        static_pos_arg = 0
        deferred_pos_arg = beam.pvalue.AsSingleton(
            p | "CreateDeferredSideInput" >> beam.Create([1]))
        static_kwarg = 2
        deferred_kwarg = beam.pvalue.AsSingleton(
            p | "CreateDeferredSideInputKwarg" >> beam.Create([3]))
        res = (
            p
            | "CreateInputs" >> beam.Create([(None, None)])
            | beam.CombinePerKey(
                MyCombiner(),
                static_pos_arg,
                deferred_pos_arg,
                static_kwarg=static_kwarg,
                deferred_kwarg=deferred_kwarg))
        assert_that(
            res,
            equal_to([
                (None, ((0, 1), {
                    'static_kwarg': 2, 'deferred_kwarg': 3
                }))
            ]))

      # Check that the combinefn was called with the expected arguments
      with open(fname, "r") as f:
        data = json.load(f)
        expected_args = [0, 1]
        expected_kwargs = {"static_kwarg": 2, "deferred_kwarg": 3}
        method_names = [
            "create_accumulator",
            "compact",
            "add_input",
            "merge_accumulators",
            "extract_output"
        ]
        for key in method_names:
          print(f"Checking {key}")
          self.assertEqual(data[key + "_args"], expected_args)
          self.assertEqual(data[key + "_kwargs"], expected_kwargs)

  def test_cpk_with_windows(self):
    # With global window side input
    with TestPipeline() as p:

      def sum_with_floor(vals, min_value=0):
        vals_sum = sum(vals)
        if vals_sum < min_value:
          vals_sum += min_value
        return vals_sum

      res = (
          p
          | "CreateInputs" >> beam.Create([1, 2, 100, 101, 102])
          | beam.Map(lambda x: window.TimestampedValue(('k', x), x))
          | beam.WindowInto(FixedWindows(99))
          | beam.CombinePerKey(
              sum_with_floor,
              min_value=pvalue.AsSingleton(p | beam.Create([100]))))
      assert_that(res, equal_to([('k', 103), ('k', 303)]))

    # with matching window side input
    with TestPipeline() as p:
      min_value = (
          p
          | "CreateMinValue" >> beam.Create([
              window.TimestampedValue(50, 5),
              window.TimestampedValue(1000, 100)
          ])
          | "WindowSideInputs" >> beam.WindowInto(FixedWindows(99)))
      res = (
          p
          | "CreateInputs" >> beam.Create([1, 2, 100, 101, 102])
          | beam.Map(lambda x: window.TimestampedValue(('k', x), x))
          | beam.WindowInto(FixedWindows(99))
          | beam.CombinePerKey(
              sum_with_floor, min_value=pvalue.AsSingleton(min_value)))
      assert_that(res, equal_to([('k', 53), ('k', 1303)]))


if __name__ == '__main__':
  unittest.main()