/*
 * 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.
 */
package org.apache.beam.runners.flink.translation.wrappers.streaming.io.source;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.apache.beam.runners.flink.FlinkPipelineOptions;
import org.apache.beam.runners.flink.translation.wrappers.streaming.io.TestBoundedCountingSource;
import org.apache.beam.runners.flink.translation.wrappers.streaming.io.TestCountingSource;
import org.apache.beam.sdk.io.Source;
import org.apache.beam.sdk.values.KV;
import org.apache.flink.connector.testutils.source.reader.TestingSplitEnumeratorContext;
import org.junit.Test;

/** Unit tests for {@link FlinkSourceSplitEnumerator}. */
public class FlinkSourceSplitEnumeratorTest {

  @Test
  public void testAssignSplitsWithBoundedSource() throws IOException {
    final int numSubtasks = 2;
    final int numSplits = 10;
    final int totalNumRecords = 10;
    TestingSplitEnumeratorContext<FlinkSourceSplit<KV<Integer, Integer>>> testContext =
        new TestingSplitEnumeratorContext<>(numSubtasks);
    TestBoundedCountingSource testSource =
        new TestBoundedCountingSource(numSplits, totalNumRecords);

    assignSplits(testContext, testSource, numSplits);
    assertEquals(numSubtasks, testContext.getSplitAssignments().size());

    testContext
        .getSplitAssignments()
        .forEach(
            (subtaskId, state) -> {
              int expectedNumSplitsPerSubtask = numSplits / numSubtasks;
              assertEquals(
                  "Each subtask should have " + expectedNumSplitsPerSubtask + " assigned splits",
                  expectedNumSplitsPerSubtask,
                  state.getAssignedSplits().size());
              assertTrue(
                  "Each subtask should have received NoMoreSplits",
                  state.hasReceivedNoMoreSplitsSignal());
              state
                  .getAssignedSplits()
                  .forEach(
                      split -> {
                        TestBoundedCountingSource source =
                            (TestBoundedCountingSource) split.getBeamSplitSource();
                        try {
                          int expectedSplitSize = totalNumRecords / numSplits;
                          assertEquals(
                              expectedSplitSize,
                              source.getEstimatedSizeBytes(FlinkPipelineOptions.defaults()));
                        } catch (Exception e) {
                          fail("Received exception" + e);
                        }
                      });
            });
  }

  @Test
  public void testAssignSplitsWithUnboundedSource() throws IOException {
    final int numSplits = 10;
    final int numSubtasks = 5;
    final int numRecordsPerSplit = 10;
    TestingSplitEnumeratorContext<FlinkSourceSplit<KV<Integer, Integer>>> testContext =
        new TestingSplitEnumeratorContext<>(numSubtasks);
    TestCountingSource testSource = new TestCountingSource(numRecordsPerSplit);

    assignSplits(testContext, testSource, numSplits);

    testContext
        .getSplitAssignments()
        .forEach(
            (subtaskId, state) -> {
              int expectedNumSplitsPerSubtask = numSplits / numSubtasks;
              assertEquals(
                  "Each subtask should have " + expectedNumSplitsPerSubtask + " assigned splits",
                  expectedNumSplitsPerSubtask,
                  state.getAssignedSplits().size());
              assertTrue(
                  "Each subtask should have received NoMoreSplits",
                  state.hasReceivedNoMoreSplitsSignal());
            });
  }

  @Test
  public void testAddSplitsBack() throws IOException {
    final int numSubtasks = 2;
    final int numSplits = 10;
    final int totalNumRecords = 10;
    TestingSplitEnumeratorContext<FlinkSourceSplit<KV<Integer, Integer>>> testContext =
        new TestingSplitEnumeratorContext<>(numSubtasks);
    TestBoundedCountingSource testSource =
        new TestBoundedCountingSource(numSplits, totalNumRecords);
    try (FlinkSourceSplitEnumerator<KV<Integer, Integer>> splitEnumerator =
        new FlinkSourceSplitEnumerator<>(
            testContext, testSource, FlinkPipelineOptions.defaults(), numSplits)) {
      splitEnumerator.start();
      testContext.registerReader(0, "0");
      splitEnumerator.addReader(0);
      testContext.getExecutorService().triggerAll();

      List<FlinkSourceSplit<KV<Integer, Integer>>> splitsForReader =
          testContext.getSplitAssignments().get(0).getAssignedSplits();
      assertEquals(numSplits / numSubtasks, splitsForReader.size());

      splitEnumerator.addSplitsBack(splitsForReader, 0);
      splitEnumerator.addReader(0);
      assertEquals(2 * numSplits / numSubtasks, splitsForReader.size());
    }
  }

  @Test
  public void testAddSplitsBackAfterRescale() throws Exception {
    final int numSubtasks = 2;
    final int numSplits = 10;
    final int totalNumRecords = 10;
    TestingSplitEnumeratorContext<FlinkSourceSplit<KV<Integer, Integer>>> testContext =
        new TestingSplitEnumeratorContext<>(numSubtasks);
    TestBoundedCountingSource testSource =
        new TestBoundedCountingSource(numSplits, totalNumRecords);
    final Map<Integer, List<FlinkSourceSplit<KV<Integer, Integer>>>> assignment;
    try (FlinkSourceSplitEnumerator<KV<Integer, Integer>> splitEnumerator =
        new FlinkSourceSplitEnumerator<>(
            testContext, testSource, FlinkPipelineOptions.defaults(), numSplits)) {
      splitEnumerator.start();
      for (int i = 0; i < numSubtasks; i++) {
        testContext.registerReader(i, String.valueOf(i));
        splitEnumerator.addReader(i);
      }
      testContext.getExecutorService().triggerAll();
      assignment =
          testContext.getSplitAssignments().entrySet().stream()
              .map(e -> KV.of(e.getKey(), e.getValue().getAssignedSplits()))
              .collect(Collectors.toMap(KV::getKey, KV::getValue));
    }

    // add tasks back
    testContext = new TestingSplitEnumeratorContext<>(numSubtasks);
    try (FlinkSourceSplitEnumerator<KV<Integer, Integer>> splitEnumerator =
        new FlinkSourceSplitEnumerator<>(
            testContext, testSource, FlinkPipelineOptions.defaults(), numSplits, true)) {
      splitEnumerator.start();
      assignment.forEach(
          (splitId, assignedSplits) -> splitEnumerator.addSplitsBack(assignedSplits, splitId));
      testContext.registerReader(0, "0");
      splitEnumerator.addReader(0);
      testContext.getExecutorService().triggerAll();

      List<FlinkSourceSplit<KV<Integer, Integer>>> splitsForReader =
          testContext.getSplitAssignments().get(0).getAssignedSplits();
      assertEquals(numSplits / numSubtasks, splitsForReader.size());
    }
  }

  private void assignSplits(
      TestingSplitEnumeratorContext<FlinkSourceSplit<KV<Integer, Integer>>> context,
      Source<KV<Integer, Integer>> source,
      int numSplits)
      throws IOException {
    try (FlinkSourceSplitEnumerator<KV<Integer, Integer>> splitEnumerator =
        new FlinkSourceSplitEnumerator<>(
            context, source, FlinkPipelineOptions.defaults(), numSplits)) {
      splitEnumerator.start();
      // Add a reader before splitting the beam source.
      context.registerReader(0, "0");
      splitEnumerator.addReader(0);
      context.getExecutorService().triggerAll();
      context.registerReader(1, "1");
      // Add another reader after splitting the beam source.
      splitEnumerator.addReader(1);
    }
  }
}
