/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/core/kernels/data/shuffle_dataset_op.h"

#include "tensorflow/core/kernels/data/dataset_test_base.h"

namespace tensorflow {
namespace data {
namespace {

constexpr char kShuffleNodeName[] = "shuffle_dataset";
constexpr char kShuffleAndRepeatNodeName[] = "shuffle_and_repeat_dataset";

class ShuffleDatasetOpTest : public DatasetOpsTestBase {
 protected:
  // Creates a new `ShuffleDataset`/`ShuffleAndRepeatDataset` op kernel
  Status CreateDatasetOpKernel(
      int64 count, bool reshuffle_each_iteration,
      const DataTypeVector& output_types,
      const std::vector<PartialTensorShape>& output_shapes,
      std::unique_ptr<OpKernel>* shuffle_dataset_kernel) {
    NodeDef node_def;
    if (count == 1) {
      node_def = test::function::NDef(
          kShuffleNodeName, name_utils::OpName(ShuffleDatasetOp::kDatasetType),
          {ShuffleDatasetOp::kInputDataset, ShuffleDatasetOp::kBufferSize,
           ShuffleDatasetOp::kSeed, ShuffleDatasetOp::kSeed2},
          {{ShuffleDatasetOp::kReshuffleEachIteration,
            reshuffle_each_iteration},
           {ShuffleDatasetOp::kOutputTypes, output_types},
           {ShuffleDatasetOp::kOutputShapes, output_shapes}});
    } else {
      node_def = test::function::NDef(
          kShuffleAndRepeatNodeName,
          name_utils::OpName(ShuffleAndRepeatDatasetOp::kDatasetType),
          {ShuffleAndRepeatDatasetOp::kInputDataset,
           ShuffleAndRepeatDatasetOp::kBufferSize,
           ShuffleAndRepeatDatasetOp::kSeed, ShuffleAndRepeatDatasetOp::kSeed2,
           ShuffleAndRepeatDatasetOp::kCount},
          {{ShuffleAndRepeatDatasetOp::kOutputTypes, output_types},
           {ShuffleAndRepeatDatasetOp::kOutputShapes, output_shapes}});
    }
    TF_RETURN_IF_ERROR(CreateOpKernel(node_def, shuffle_dataset_kernel));
    return Status::OK();
  }

  // Creates a new `ShuffleDataset`/`ShuffleAndRepeatDataset` op kernel context.
  Status CreateDatasetContext(OpKernel* const op_kernel,
                              gtl::InlinedVector<TensorValue, 4>* const inputs,
                              std::unique_ptr<OpKernelContext>* context) {
    TF_RETURN_IF_ERROR(CheckOpKernelInput(*op_kernel, *inputs));
    TF_RETURN_IF_ERROR(CreateOpKernelContext(op_kernel, inputs, context));
    return Status::OK();
  }
};

struct RangeDatasetParam {
  int64 start;
  int64 end;
  int64 step;
};

struct TestCase {
  RangeDatasetParam range_data_param;
  Tensor buffer_size;
  Tensor seed;
  Tensor seed2;
  Tensor count;
  bool reshuffle_each_iteration;
  std::vector<Tensor> expected_shuffle_outputs;
  std::vector<Tensor> expected_reshuffle_outputs;
  DataTypeVector expected_output_dtypes;
  std::vector<PartialTensorShape> expected_output_shapes;
  int64 expected_cardinality;
  std::vector<int> breakpoints;
};

template <typename T>
std::vector<Tensor> ConvertToTensorVec(std::vector<T> values) {
  std::vector<Tensor> tensors;
  tensors.reserve(values.size());
  for (auto& value : values) {
    tensors.emplace_back(CreateTensor<T>(TensorShape({}), {value}));
  }
  return tensors;
}

// Test case 1: test shuffle_dataset with reshuffle_each_iteration = false.
TestCase TestCase1() {
  return {/*range_data_param*/ {0, 10, 1},
          /*buffer_size*/
          CreateTensor<int64>(TensorShape({}), {3}),
          /*seed*/ CreateTensor<int64>(TensorShape({}), {1}),
          /*seed2*/ CreateTensor<int64>(TensorShape({}), {2}),
          /*count*/ CreateTensor<int64>(TensorShape({}), {1}),
          /*reshuffle_each_iteration*/ false,
          /*expected_shuffle_outputs*/
          ConvertToTensorVec<int64>({2, 3, 0, 5, 6, 4, 7, 8, 9, 1}),
          /*expected_reshuffle_outputs*/
          ConvertToTensorVec<int64>({2, 3, 0, 5, 6, 4, 7, 8, 9, 1}),
          /*expected_output_dtypes*/ {DT_INT64},
          /*expected_output_shapes*/ {PartialTensorShape({})},
          /*expected_cardinality*/ 10,
          /*breakpoints*/ {0, 1, 9}};
}

// Test case 2: test shuffle_dataset with reshuffle_each_iteration = true.
TestCase TestCase2() {
  return {/*range_data_param*/ {0, 10, 1},
          /*buffer_size*/
          CreateTensor<int64>(TensorShape({}), {10}),
          /*seed*/ CreateTensor<int64>(TensorShape({}), {1}),
          /*seed2*/ CreateTensor<int64>(TensorShape({}), {2}),
          /*count*/ CreateTensor<int64>(TensorShape({}), {1}),
          /*reshuffle_each_iteration*/ true,
          /*expected_shuffle_outputs*/
          ConvertToTensorVec<int64>({2, 6, 1, 3, 9, 5, 0, 8, 7, 4}),
          /*expected_reshuffle_outputs*/
          ConvertToTensorVec<int64>({1, 6, 0, 5, 2, 7, 4, 3, 9, 8}),
          /*expected_output_dtypes*/ {DT_INT64},
          /*expected_output_shapes*/ {PartialTensorShape({})},
          /*expected_cardinality*/ 10,
          /*breakpoints*/ {0, 1, 9}};
}

// Test case 3: similar with the test case 2 but a smaller buffer size than
// the input dataset.
TestCase TestCase3() {
  return {/*range_data_param*/ {0, 10, 1},
          /*buffer_size*/
          CreateTensor<int64>(TensorShape({}), {2}),
          /*seed*/ CreateTensor<int64>(TensorShape({}), {1}),
          /*seed2*/ CreateTensor<int64>(TensorShape({}), {2}),
          /*count*/ CreateTensor<int64>(TensorShape({}), {1}),
          /*reshuffle_each_iteration*/ true,
          /*expected_shuffle_outputs*/
          ConvertToTensorVec<int64>({0, 2, 1, 3, 5, 6, 4, 7, 8, 9}),
          /*expected_reshuffle_outputs*/
          ConvertToTensorVec<int64>({1, 0, 2, 3, 4, 5, 6, 7, 9, 8}),
          /*expected_output_dtypes*/ {DT_INT64},
          /*expected_output_shapes*/ {PartialTensorShape({})},
          /*expected_cardinality*/ 10,
          /*breakpoints*/ {0, 1, 9}};
}

// Test case 4: similar with the test case 2 but has different seeds.
TestCase TestCase4() {
  return {/*range_data_param*/ {0, 10, 1},
          /*buffer_size*/
          CreateTensor<int64>(TensorShape({}), {10}),
          /*seed*/ CreateTensor<int64>(TensorShape({}), {2}),
          /*seed2*/ CreateTensor<int64>(TensorShape({}), {2}),
          /*count*/ CreateTensor<int64>(TensorShape({}), {1}),
          /*reshuffle_each_iteration*/ true,
          /*expected_shuffle_outputs*/
          ConvertToTensorVec<int64>({3, 0, 8, 1, 5, 4, 7, 2, 6, 9}),
          /*expected_reshuffle_outputs*/
          ConvertToTensorVec<int64>({4, 6, 9, 0, 1, 8, 2, 7, 3, 5}),
          /*expected_output_dtypes*/ {DT_INT64},
          /*expected_output_shapes*/ {PartialTensorShape({})},
          /*expected_cardinality*/ 10,
          /*breakpoints*/ {0, 1, 9}};
}

// Test case 5: test shuffle_dataset with buffer_size = 1 &
// reshuffle_each_iteration = true.
TestCase TestCase5() {
  return {/*range_data_param*/ {0, 10, 1},
          /*buffer_size*/
          CreateTensor<int64>(TensorShape({}), {1}),
          /*seed*/ CreateTensor<int64>(TensorShape({}), {1}),
          /*seed2*/ CreateTensor<int64>(TensorShape({}), {2}),
          /*count*/ CreateTensor<int64>(TensorShape({}), {1}),
          /*reshuffle_each_iteration*/ true,
          /*expected_shuffle_outputs*/
          ConvertToTensorVec<int64>({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}),
          /*expected_reshuffle_outputs*/
          ConvertToTensorVec<int64>({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}),
          /*expected_output_dtypes*/ {DT_INT64},
          /*expected_output_shapes*/ {PartialTensorShape({})},
          /*expected_cardinality*/ 10,
          /*breakpoints*/ {0, 1, 9}};
}

// Test case 6: test shuffle_dataset with an empty input dataset.
TestCase TestCase6() {
  return {/*range_data_param*/ {0, 0, 1},
          /*buffer_size*/
          CreateTensor<int64>(TensorShape({}), {10}),
          /*seed*/ CreateTensor<int64>(TensorShape({}), {1}),
          /*seed2*/ CreateTensor<int64>(TensorShape({}), {2}),
          /*count*/ CreateTensor<int64>(TensorShape({}), {1}),
          /*reshuffle_each_iteration*/ true,
          /*expected_shuffle_outputs*/
          ConvertToTensorVec<int64>({}),
          /*expected_reshuffle_outputs*/
          ConvertToTensorVec<int64>({}),
          /*expected_output_dtypes*/ {DT_INT64},
          /*expected_output_shapes*/ {PartialTensorShape({})},
          /*expected_cardinality*/ 0,
          /*breakpoints*/ {0, 1, 9}};
}

// Test case 7: test shuffle_and_repeat_dataset with buffer_size = 10 &
// count = 2.
TestCase TestCase7() {
  return {/*range_data_param*/ {0, 10, 1},
          /*buffer_size*/
          CreateTensor<int64>(TensorShape({}), {10}),
          /*seed*/ CreateTensor<int64>(TensorShape({}), {1}),
          /*seed2*/ CreateTensor<int64>(TensorShape({}), {2}),
          /*count*/ CreateTensor<int64>(TensorShape({}), {2}),
          /*reshuffle_each_iteration*/ false,
          /*expected_shuffle_outputs*/
          ConvertToTensorVec<int64>(
              {9, 0, 8, 6, 1, 3, 7, 2, 4, 5, 4, 3, 0, 5, 8, 2, 6, 9, 7, 1}),
          /*expected_reshuffle_outputs*/
          ConvertToTensorVec<int64>(
              {9, 0, 8, 6, 1, 3, 7, 2, 4, 5, 4, 3, 0, 5, 8, 2, 6, 9, 7, 1}),
          /*expected_output_dtypes*/ {DT_INT64},
          /*expected_output_shapes*/ {PartialTensorShape({})},
          /*expected_cardinality*/ 20,
          /*breakpoints*/ {0, 5, 22}};
}

// Test case 8: test shuffle_and_repeat_dataset with buffer_size = 10 &
// count = -1
TestCase TestCase8() {
  return {/*range_data_param*/ {0, 3, 1},
          /*buffer_size*/
          CreateTensor<int64>(TensorShape({}), {10}),
          /*seed*/ CreateTensor<int64>(TensorShape({}), {1}),
          /*seed2*/ CreateTensor<int64>(TensorShape({}), {2}),
          /*count*/ CreateTensor<int64>(TensorShape({}), {-1}),
          /*reshuffle_each_iteration*/ false,
          /*expected_shuffle_outputs*/
          ConvertToTensorVec<int64>(
              {2, 0, 1, 2, 0, 1, 1, 2, 0, 1, 0, 2, 2, 0, 1, 1, 0, 2, 2, 1, 0}),
          /*expected_reshuffle_outputs*/
          ConvertToTensorVec<int64>(
              {2, 0, 1, 2, 0, 1, 1, 2, 0, 1, 0, 2, 2, 0, 1, 1, 0, 2, 2, 1, 0}),
          /*expected_output_dtypes*/ {DT_INT64},
          /*expected_output_shapes*/ {PartialTensorShape({})},
          /*expected_cardinality*/ kInfiniteCardinality,
          /*breakpoints*/ {0, 5, 20}};
}

TestCase InvalidBufferSizeTestCaseForShuffleDataset() {
  return {/*range_data_param*/ {0, 10, 1},
          /*buffer_size*/
          CreateTensor<int64>(TensorShape({}), {-1}),
          /*seed*/ CreateTensor<int64>(TensorShape({}), {1}),
          /*seed2*/ CreateTensor<int64>(TensorShape({}), {2}),
          /*count*/ CreateTensor<int64>(TensorShape({}), {1}),
          /*reshuffle_each_iteration*/ true,
          /*expected_shuffle_outputs*/ ConvertToTensorVec<int64>({}),
          /*expected_reshuffle_outputs*/ ConvertToTensorVec<int64>({}),
          /*expected_output_dtypes*/ {DT_INT64},
          /*expected_output_shapes*/ {PartialTensorShape({})},
          /*expected_cardinality*/ 0,
          /*breakpoints*/ {0, 1, 9}};
}

TestCase InvalidBufferSizeTestCaseForShuffleAndRepeatDataset() {
  return {/*range_data_param*/ {0, 10, 1},
          /*buffer_size*/
          CreateTensor<int64>(TensorShape({}), {-1}),
          /*seed*/ CreateTensor<int64>(TensorShape({}), {1}),
          /*seed2*/ CreateTensor<int64>(TensorShape({}), {2}),
          /*count*/ CreateTensor<int64>(TensorShape({}), {2}),
          /*reshuffle_each_iteration*/ true,
          /*expected_shuffle_outputs*/ ConvertToTensorVec<int64>({}),
          /*expected_reshuffle_outputs*/ ConvertToTensorVec<int64>({}),
          /*expected_output_dtypes*/ {DT_INT64},
          /*expected_output_shapes*/ {PartialTensorShape({})},
          /*expected_cardinality*/ 0,
          /*breakpoints*/ {0, 1, 9}};
}

TestCase InvalidCountTestCaseForShuffleAndRepeatDataset() {
  return {/*range_data_param*/ {0, 3, 1},
          /*buffer_size*/
          CreateTensor<int64>(TensorShape({}), {10}),
          /*seed*/ CreateTensor<int64>(TensorShape({}), {1}),
          /*seed2*/ CreateTensor<int64>(TensorShape({}), {2}),
          /*count*/ CreateTensor<int64>(TensorShape({}), {0}),
          /*reshuffle_each_iteration*/ false,
          /*expected_shuffle_outputs*/
          ConvertToTensorVec<int64>({}),
          /*expected_reshuffle_outputs*/
          ConvertToTensorVec<int64>({}),
          /*expected_output_dtypes*/ {DT_INT64},
          /*expected_output_shapes*/ {PartialTensorShape({})},
          /*expected_cardinality*/ 0,
          /*breakpoints*/ {0, 5, 20}};
}

class ParameterizedShuffleDatasetOpTest
    : public ShuffleDatasetOpTest,
      public ::testing::WithParamInterface<TestCase> {};

TEST_P(ParameterizedShuffleDatasetOpTest, GetNext) {
  int thread_num = 2, cpu_num = 2;
  TestCase test_case = GetParam();
  TF_ASSERT_OK(InitThreadPool(thread_num));
  TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));

  Tensor count = test_case.count;
  int64 count_value = count.flat<int64>()(0);
  std::unique_ptr<OpKernel> dataset_kernel;
  TF_ASSERT_OK(
      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
                            test_case.expected_output_dtypes,
                            test_case.expected_output_shapes, &dataset_kernel));

  DatasetBase* range_dataset;
  TF_ASSERT_OK(CreateRangeDataset<int64>(
      test_case.range_data_param.start, test_case.range_data_param.end,
      test_case.range_data_param.step, "range", &range_dataset));
  Tensor range_dataset_tensor(DT_VARIANT, TensorShape({}));
  TF_ASSERT_OK(
      StoreDatasetInVariantTensor(range_dataset, &range_dataset_tensor));
  Tensor buffer_size = test_case.buffer_size;
  Tensor seed = test_case.seed;
  Tensor seed2 = test_case.seed2;
  gtl::InlinedVector<TensorValue, 4> inputs(
      {TensorValue(&range_dataset_tensor), TensorValue(&buffer_size),
       TensorValue(&seed), TensorValue(&seed2)});
  if (count_value != 1) inputs.push_back(TensorValue(&count));

  std::unique_ptr<OpKernelContext> dataset_context;
  TF_ASSERT_OK(
      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
  DatasetBase* dataset;
  TF_ASSERT_OK(
      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
  core::ScopedUnref scoped_unref_dataset(dataset);

  std::unique_ptr<IteratorContext> iterator_ctx;
  TF_ASSERT_OK(CreateIteratorContext(dataset_context.get(), &iterator_ctx));
  std::unique_ptr<IteratorBase> iterator;
  TF_ASSERT_OK(
      dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator));

  bool end_of_sequence = false;
  std::vector<Tensor> shuffled_out_tensors;
  while (!end_of_sequence) {
    std::vector<Tensor> next;
    TF_EXPECT_OK(
        iterator->GetNext(iterator_ctx.get(), &next, &end_of_sequence));
    shuffled_out_tensors.insert(shuffled_out_tensors.end(), next.begin(),
                                next.end());
    // For the forever-repeat case, we test only a finite number of steps of
    // the infinite sequence.
    if (count_value == -1 && shuffled_out_tensors.size() ==
                                 test_case.expected_shuffle_outputs.size()) {
      break;
    }
  }

  // Reshuffle the dataset.
  end_of_sequence = false;
  TF_ASSERT_OK(
      dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator));
  std::vector<Tensor> reshuffled_out_tensors;
  while (!end_of_sequence) {
    std::vector<Tensor> next;
    TF_EXPECT_OK(
        iterator->GetNext(iterator_ctx.get(), &next, &end_of_sequence));
    reshuffled_out_tensors.insert(reshuffled_out_tensors.end(), next.begin(),
                                  next.end());
    // For the forever-repeat case, we test only a finite number of steps of
    // the infinite sequence.
    if (count_value == -1 && reshuffled_out_tensors.size() ==
                                 test_case.expected_shuffle_outputs.size()) {
      break;
    }
  }

  TF_EXPECT_OK(ExpectEqual(shuffled_out_tensors,
                           test_case.expected_shuffle_outputs,
                           /*compare_order*/ true));
  TF_EXPECT_OK(ExpectEqual(reshuffled_out_tensors,
                           test_case.expected_reshuffle_outputs,
                           /*compare_order*/ true));
}

TEST_P(ParameterizedShuffleDatasetOpTest, DatasetNodeName) {
  int thread_num = 2, cpu_num = 2;
  TestCase test_case = GetParam();
  TF_ASSERT_OK(InitThreadPool(thread_num));
  TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));

  Tensor count = test_case.count;
  int64 count_value = count.flat<int64>()(0);
  std::unique_ptr<OpKernel> dataset_kernel;
  TF_ASSERT_OK(
      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
                            test_case.expected_output_dtypes,
                            test_case.expected_output_shapes, &dataset_kernel));

  DatasetBase* range_dataset;
  TF_ASSERT_OK(CreateRangeDataset<int64>(
      test_case.range_data_param.start, test_case.range_data_param.end,
      test_case.range_data_param.step, "range", &range_dataset));
  Tensor range_dataset_tensor(DT_VARIANT, TensorShape({}));
  TF_ASSERT_OK(
      StoreDatasetInVariantTensor(range_dataset, &range_dataset_tensor));
  Tensor buffer_size = test_case.buffer_size;
  Tensor seed = test_case.seed;
  Tensor seed2 = test_case.seed2;
  gtl::InlinedVector<TensorValue, 4> inputs(
      {TensorValue(&range_dataset_tensor), TensorValue(&buffer_size),
       TensorValue(&seed), TensorValue(&seed2)});
  if (count_value != 1) inputs.push_back(TensorValue(&count));

  std::unique_ptr<OpKernelContext> dataset_context;
  TF_ASSERT_OK(
      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
  DatasetBase* dataset;
  TF_ASSERT_OK(
      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
  core::ScopedUnref scoped_unref_dataset(dataset);

  if (count_value == 1) {
    EXPECT_EQ(dataset->node_name(), kShuffleNodeName);
  } else {
    EXPECT_EQ(dataset->node_name(), kShuffleAndRepeatNodeName);
  }
}

TEST_P(ParameterizedShuffleDatasetOpTest, DatasetTypeString) {
  int thread_num = 2, cpu_num = 2;
  TestCase test_case = GetParam();
  TF_ASSERT_OK(InitThreadPool(thread_num));
  TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));

  Tensor count = test_case.count;
  int64 count_value = count.flat<int64>()(0);
  std::unique_ptr<OpKernel> dataset_kernel;
  TF_ASSERT_OK(
      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
                            test_case.expected_output_dtypes,
                            test_case.expected_output_shapes, &dataset_kernel));

  DatasetBase* range_dataset;
  TF_ASSERT_OK(CreateRangeDataset<int64>(
      test_case.range_data_param.start, test_case.range_data_param.end,
      test_case.range_data_param.step, "range", &range_dataset));
  Tensor range_dataset_tensor(DT_VARIANT, TensorShape({}));
  TF_ASSERT_OK(
      StoreDatasetInVariantTensor(range_dataset, &range_dataset_tensor));
  Tensor buffer_size = test_case.buffer_size;
  Tensor seed = test_case.seed;
  Tensor seed2 = test_case.seed2;
  gtl::InlinedVector<TensorValue, 4> inputs(
      {TensorValue(&range_dataset_tensor), TensorValue(&buffer_size),
       TensorValue(&seed), TensorValue(&seed2)});
  if (count_value != 1) inputs.push_back(TensorValue(&count));

  std::unique_ptr<OpKernelContext> dataset_context;
  TF_ASSERT_OK(
      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
  DatasetBase* dataset;
  TF_ASSERT_OK(
      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
  core::ScopedUnref scoped_unref_dataset(dataset);

  if (count_value == 1) {
    EXPECT_EQ(dataset->type_string(),
              name_utils::OpName(ShuffleDatasetOp::kDatasetType));
  } else {
    EXPECT_EQ(dataset->type_string(),
              name_utils::OpName(ShuffleAndRepeatDatasetOp::kDatasetType));
  }
}

TEST_P(ParameterizedShuffleDatasetOpTest, DatasetOutputDtypes) {
  int thread_num = 2, cpu_num = 2;
  TestCase test_case = GetParam();
  TF_ASSERT_OK(InitThreadPool(thread_num));
  TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));

  Tensor count = test_case.count;
  int64 count_value = count.flat<int64>()(0);
  std::unique_ptr<OpKernel> dataset_kernel;
  TF_ASSERT_OK(
      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
                            test_case.expected_output_dtypes,
                            test_case.expected_output_shapes, &dataset_kernel));

  DatasetBase* range_dataset;
  TF_ASSERT_OK(CreateRangeDataset<int64>(
      test_case.range_data_param.start, test_case.range_data_param.end,
      test_case.range_data_param.step, "range", &range_dataset));
  Tensor range_dataset_tensor(DT_VARIANT, TensorShape({}));
  TF_ASSERT_OK(
      StoreDatasetInVariantTensor(range_dataset, &range_dataset_tensor));
  Tensor buffer_size = test_case.buffer_size;
  Tensor seed = test_case.seed;
  Tensor seed2 = test_case.seed2;
  gtl::InlinedVector<TensorValue, 4> inputs(
      {TensorValue(&range_dataset_tensor), TensorValue(&buffer_size),
       TensorValue(&seed), TensorValue(&seed2)});
  if (count_value != 1) inputs.push_back(TensorValue(&count));

  std::unique_ptr<OpKernelContext> dataset_context;
  TF_ASSERT_OK(
      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
  DatasetBase* dataset;
  TF_ASSERT_OK(
      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
  core::ScopedUnref scoped_unref_dataset(dataset);

  TF_EXPECT_OK(VerifyTypesMatch(dataset->output_dtypes(),
                                test_case.expected_output_dtypes));
}

TEST_P(ParameterizedShuffleDatasetOpTest, DatasetOutputShapes) {
  int thread_num = 2, cpu_num = 2;
  TestCase test_case = GetParam();
  TF_ASSERT_OK(InitThreadPool(thread_num));
  TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));

  Tensor count = test_case.count;
  int64 count_value = count.flat<int64>()(0);
  std::unique_ptr<OpKernel> dataset_kernel;
  TF_ASSERT_OK(
      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
                            test_case.expected_output_dtypes,
                            test_case.expected_output_shapes, &dataset_kernel));

  DatasetBase* range_dataset;
  TF_ASSERT_OK(CreateRangeDataset<int64>(
      test_case.range_data_param.start, test_case.range_data_param.end,
      test_case.range_data_param.step, "range", &range_dataset));
  Tensor range_dataset_tensor(DT_VARIANT, TensorShape({}));
  TF_ASSERT_OK(
      StoreDatasetInVariantTensor(range_dataset, &range_dataset_tensor));
  Tensor buffer_size = test_case.buffer_size;
  Tensor seed = test_case.seed;
  Tensor seed2 = test_case.seed2;
  gtl::InlinedVector<TensorValue, 4> inputs(
      {TensorValue(&range_dataset_tensor), TensorValue(&buffer_size),
       TensorValue(&seed), TensorValue(&seed2)});
  if (count_value != 1) inputs.push_back(TensorValue(&count));

  std::unique_ptr<OpKernelContext> dataset_context;
  TF_ASSERT_OK(
      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
  DatasetBase* dataset;
  TF_ASSERT_OK(
      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
  core::ScopedUnref scoped_unref_dataset(dataset);

  TF_EXPECT_OK(VerifyShapesCompatible(dataset->output_shapes(),
                                      test_case.expected_output_shapes));
}

TEST_P(ParameterizedShuffleDatasetOpTest, Cardinality) {
  int thread_num = 2, cpu_num = 2;
  TestCase test_case = GetParam();
  TF_ASSERT_OK(InitThreadPool(thread_num));
  TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));

  Tensor count = test_case.count;
  int64 count_value = count.flat<int64>()(0);
  std::unique_ptr<OpKernel> dataset_kernel;
  TF_ASSERT_OK(
      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
                            test_case.expected_output_dtypes,
                            test_case.expected_output_shapes, &dataset_kernel));

  DatasetBase* range_dataset;
  TF_ASSERT_OK(CreateRangeDataset<int64>(
      test_case.range_data_param.start, test_case.range_data_param.end,
      test_case.range_data_param.step, "range", &range_dataset));
  Tensor range_dataset_tensor(DT_VARIANT, TensorShape({}));
  TF_ASSERT_OK(
      StoreDatasetInVariantTensor(range_dataset, &range_dataset_tensor));
  Tensor buffer_size = test_case.buffer_size;
  Tensor seed = test_case.seed;
  Tensor seed2 = test_case.seed2;
  gtl::InlinedVector<TensorValue, 4> inputs(
      {TensorValue(&range_dataset_tensor), TensorValue(&buffer_size),
       TensorValue(&seed), TensorValue(&seed2)});
  if (count_value != 1) inputs.push_back(TensorValue(&count));

  std::unique_ptr<OpKernelContext> dataset_context;
  TF_ASSERT_OK(
      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
  DatasetBase* dataset;
  TF_ASSERT_OK(
      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
  core::ScopedUnref scoped_unref_dataset(dataset);

  EXPECT_EQ(dataset->Cardinality(), test_case.expected_cardinality);
}

TEST_P(ParameterizedShuffleDatasetOpTest, IteratorOutputDtypes) {
  int thread_num = 2, cpu_num = 2;
  TestCase test_case = GetParam();
  TF_ASSERT_OK(InitThreadPool(thread_num));
  TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));

  Tensor count = test_case.count;
  int64 count_value = count.flat<int64>()(0);
  std::unique_ptr<OpKernel> dataset_kernel;
  TF_ASSERT_OK(
      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
                            test_case.expected_output_dtypes,
                            test_case.expected_output_shapes, &dataset_kernel));

  DatasetBase* range_dataset;
  TF_ASSERT_OK(CreateRangeDataset<int64>(
      test_case.range_data_param.start, test_case.range_data_param.end,
      test_case.range_data_param.step, "range", &range_dataset));
  Tensor range_dataset_tensor(DT_VARIANT, TensorShape({}));
  TF_ASSERT_OK(
      StoreDatasetInVariantTensor(range_dataset, &range_dataset_tensor));
  Tensor buffer_size = test_case.buffer_size;
  Tensor seed = test_case.seed;
  Tensor seed2 = test_case.seed2;
  gtl::InlinedVector<TensorValue, 4> inputs(
      {TensorValue(&range_dataset_tensor), TensorValue(&buffer_size),
       TensorValue(&seed), TensorValue(&seed2)});
  if (count_value != 1) inputs.push_back(TensorValue(&count));

  std::unique_ptr<OpKernelContext> dataset_context;
  TF_ASSERT_OK(
      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
  DatasetBase* dataset;
  TF_ASSERT_OK(
      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
  core::ScopedUnref scoped_unref_dataset(dataset);

  std::unique_ptr<IteratorContext> iterator_ctx;
  TF_ASSERT_OK(CreateIteratorContext(dataset_context.get(), &iterator_ctx));
  std::unique_ptr<IteratorBase> iterator;
  TF_ASSERT_OK(
      dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator));

  TF_EXPECT_OK(VerifyTypesMatch(iterator->output_dtypes(),
                                test_case.expected_output_dtypes));
}

TEST_P(ParameterizedShuffleDatasetOpTest, IteratorOutputShapes) {
  int thread_num = 2, cpu_num = 2;
  TestCase test_case = GetParam();
  TF_ASSERT_OK(InitThreadPool(thread_num));
  TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));

  Tensor count = test_case.count;
  int64 count_value = count.flat<int64>()(0);
  std::unique_ptr<OpKernel> dataset_kernel;
  TF_ASSERT_OK(
      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
                            test_case.expected_output_dtypes,
                            test_case.expected_output_shapes, &dataset_kernel));

  DatasetBase* range_dataset;
  TF_ASSERT_OK(CreateRangeDataset<int64>(
      test_case.range_data_param.start, test_case.range_data_param.end,
      test_case.range_data_param.step, "range", &range_dataset));
  Tensor range_dataset_tensor(DT_VARIANT, TensorShape({}));
  TF_ASSERT_OK(
      StoreDatasetInVariantTensor(range_dataset, &range_dataset_tensor));
  Tensor buffer_size = test_case.buffer_size;
  Tensor seed = test_case.seed;
  Tensor seed2 = test_case.seed2;
  gtl::InlinedVector<TensorValue, 4> inputs(
      {TensorValue(&range_dataset_tensor), TensorValue(&buffer_size),
       TensorValue(&seed), TensorValue(&seed2)});
  if (count_value != 1) inputs.push_back(TensorValue(&count));

  std::unique_ptr<OpKernelContext> dataset_context;
  TF_ASSERT_OK(
      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
  DatasetBase* dataset;
  TF_ASSERT_OK(
      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
  core::ScopedUnref scoped_unref_dataset(dataset);

  std::unique_ptr<IteratorContext> iterator_ctx;
  TF_ASSERT_OK(CreateIteratorContext(dataset_context.get(), &iterator_ctx));
  std::unique_ptr<IteratorBase> iterator;
  TF_ASSERT_OK(
      dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator));

  TF_EXPECT_OK(VerifyShapesCompatible(iterator->output_shapes(),
                                      test_case.expected_output_shapes));
}

TEST_P(ParameterizedShuffleDatasetOpTest, IteratorOutputPrefix) {
  int thread_num = 2, cpu_num = 2;
  TestCase test_case = GetParam();
  TF_ASSERT_OK(InitThreadPool(thread_num));
  TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));

  Tensor count = test_case.count;
  int64 count_value = count.flat<int64>()(0);
  std::unique_ptr<OpKernel> dataset_kernel;
  TF_ASSERT_OK(
      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
                            test_case.expected_output_dtypes,
                            test_case.expected_output_shapes, &dataset_kernel));

  DatasetBase* range_dataset;
  TF_ASSERT_OK(CreateRangeDataset<int64>(
      test_case.range_data_param.start, test_case.range_data_param.end,
      test_case.range_data_param.step, "range", &range_dataset));
  Tensor range_dataset_tensor(DT_VARIANT, TensorShape({}));
  TF_ASSERT_OK(
      StoreDatasetInVariantTensor(range_dataset, &range_dataset_tensor));
  Tensor buffer_size = test_case.buffer_size;
  Tensor seed = test_case.seed;
  Tensor seed2 = test_case.seed2;
  gtl::InlinedVector<TensorValue, 4> inputs(
      {TensorValue(&range_dataset_tensor), TensorValue(&buffer_size),
       TensorValue(&seed), TensorValue(&seed2)});
  if (count_value != 1) inputs.push_back(TensorValue(&count));

  std::unique_ptr<OpKernelContext> dataset_context;
  TF_ASSERT_OK(
      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
  DatasetBase* dataset;
  TF_ASSERT_OK(
      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
  core::ScopedUnref scoped_unref_dataset(dataset);

  std::unique_ptr<IteratorContext> iterator_ctx;
  TF_ASSERT_OK(CreateIteratorContext(dataset_context.get(), &iterator_ctx));
  std::unique_ptr<IteratorBase> iterator;
  TF_ASSERT_OK(
      dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator));

  if (count_value == 1) {
    EXPECT_EQ(
        iterator->prefix(),
        name_utils::IteratorPrefix(ShuffleDatasetOp::kDatasetType, "Iterator"));
  } else {
    EXPECT_EQ(iterator->prefix(),
              name_utils::IteratorPrefix(
                  ShuffleAndRepeatDatasetOp::kDatasetType, "Iterator"));
  }
}

TEST_P(ParameterizedShuffleDatasetOpTest, Roundtrip) {
  int thread_num = 2, cpu_num = 2;
  TestCase test_case = GetParam();
  TF_ASSERT_OK(InitThreadPool(thread_num));
  TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));

  Tensor count = test_case.count;
  int64 count_value = count.flat<int64>()(0);
  std::unique_ptr<OpKernel> dataset_kernel;
  TF_ASSERT_OK(
      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
                            test_case.expected_output_dtypes,
                            test_case.expected_output_shapes, &dataset_kernel));

  DatasetBase* range_dataset;
  TF_ASSERT_OK(CreateRangeDataset<int64>(
      test_case.range_data_param.start, test_case.range_data_param.end,
      test_case.range_data_param.step, "range", &range_dataset));
  Tensor range_dataset_tensor(DT_VARIANT, TensorShape({}));
  TF_ASSERT_OK(
      StoreDatasetInVariantTensor(range_dataset, &range_dataset_tensor));
  Tensor buffer_size = test_case.buffer_size;
  Tensor seed = test_case.seed;
  Tensor seed2 = test_case.seed2;
  gtl::InlinedVector<TensorValue, 4> inputs(
      {TensorValue(&range_dataset_tensor), TensorValue(&buffer_size),
       TensorValue(&seed), TensorValue(&seed2)});
  if (count_value != 1) inputs.push_back(TensorValue(&count));

  std::unique_ptr<OpKernelContext> dataset_context;
  TF_ASSERT_OK(
      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
  DatasetBase* dataset;
  TF_ASSERT_OK(
      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
  core::ScopedUnref scoped_unref_dataset(dataset);

  std::unique_ptr<IteratorContext> iterator_ctx;
  TF_ASSERT_OK(CreateIteratorContext(dataset_context.get(), &iterator_ctx));
  std::unique_ptr<IteratorBase> iterator;
  TF_ASSERT_OK(
      dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator));

  std::unique_ptr<SerializationContext> serialization_ctx;
  TF_ASSERT_OK(CreateSerializationContext(&serialization_ctx));

  bool end_of_sequence = false;
  std::vector<Tensor> out_tensors;
  int cur_iteration = 0;
  const std::vector<int>& breakpoints = test_case.breakpoints;
  for (int breakpoint : breakpoints) {
    VariantTensorData data;
    VariantTensorDataWriter writer(&data);
    TF_EXPECT_OK(iterator->Save(serialization_ctx.get(), &writer));
    TF_EXPECT_OK(writer.Flush());
    VariantTensorDataReader reader(&data);
    TF_EXPECT_OK(RestoreIterator(iterator_ctx.get(), &reader, "Iterator",
                                 *dataset, &iterator));

    while (cur_iteration <= breakpoint) {
      std::vector<Tensor> next;
      TF_EXPECT_OK(
          iterator->GetNext(iterator_ctx.get(), &next, &end_of_sequence));
      out_tensors.insert(out_tensors.end(), next.begin(), next.end());
      cur_iteration++;
    }
  }

  TF_EXPECT_OK(ExpectEqual(out_tensors, test_case.expected_shuffle_outputs,
                           /*compare_order*/ true));
}

INSTANTIATE_TEST_SUITE_P(ShuffleDatasetOpTest,
                         ParameterizedShuffleDatasetOpTest,
                         ::testing::ValuesIn(std::vector<TestCase>(
                             {TestCase1(), TestCase2(), TestCase3(),
                              TestCase4(), TestCase5(), TestCase6(),
                              TestCase7(), TestCase8()})));

TEST_F(ShuffleDatasetOpTest, InvalidArguments) {
  int thread_num = 2, cpu_num = 2;
  std::vector<TestCase> test_cases = {
      InvalidBufferSizeTestCaseForShuffleDataset(),
      InvalidBufferSizeTestCaseForShuffleAndRepeatDataset(),
      InvalidCountTestCaseForShuffleAndRepeatDataset()};
  TF_ASSERT_OK(InitThreadPool(thread_num));
  TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));

  for (const auto& test_case : test_cases) {
    Tensor count = test_case.count;
    int64 count_value = count.flat<int64>()(0);
    std::unique_ptr<OpKernel> dataset_kernel;
    TF_ASSERT_OK(CreateDatasetOpKernel(
        count_value, test_case.reshuffle_each_iteration,
        test_case.expected_output_dtypes, test_case.expected_output_shapes,
        &dataset_kernel));

    DatasetBase* range_dataset;
    TF_ASSERT_OK(CreateRangeDataset<int64>(
        test_case.range_data_param.start, test_case.range_data_param.end,
        test_case.range_data_param.step, "range", &range_dataset));
    Tensor range_dataset_tensor(DT_VARIANT, TensorShape({}));
    TF_ASSERT_OK(
        StoreDatasetInVariantTensor(range_dataset, &range_dataset_tensor));
    Tensor buffer_size = test_case.buffer_size;
    Tensor seed = test_case.seed;
    Tensor seed2 = test_case.seed2;
    gtl::InlinedVector<TensorValue, 4> inputs(
        {TensorValue(&range_dataset_tensor), TensorValue(&buffer_size),
         TensorValue(&seed), TensorValue(&seed2)});
    if (count_value != 1) inputs.push_back(TensorValue(&count));

    std::unique_ptr<OpKernelContext> dataset_context;
    TF_ASSERT_OK(
        CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
    DatasetBase* shuffle_dataset;
    EXPECT_EQ(CreateDataset(dataset_kernel.get(), dataset_context.get(),
                            &shuffle_dataset)
                  .code(),
              tensorflow::error::INVALID_ARGUMENT);
  }
}

}  // namespace
}  // namespace data
}  // namespace tensorflow
