Function TEST_P

tensorflow/core/kernels/data/shard_dataset_op_test.cc:262–315

      public ::testing::WithParamInterface<TestCase> {};

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

  std::unique_ptr<OpKernel> shard_dataset_kernel;
  TF_ASSERT_OK(CreateShardDatasetOpKernel(
      test_case.require_non_empty, test_case.expected_output_dtypes,
      test_case.expected_output_shapes, &shard_dataset_kernel));

  DatasetBase* range_dataset;
  TF_ASSERT_OK(CreateRangeDataset<int64>(
      test_case.range_dataset_param.start, test_case.range_dataset_param.end,
      test_case.range_dataset_param.step, "range", &range_dataset));
  Tensor range_dataset_tensor(DT_VARIANT, TensorShape({}));
  TF_ASSERT_OK(
      StoreDatasetInVariantTensor(range_dataset, &range_dataset_tensor));

  Tensor num_shards = test_case.num_shards;
  Tensor index = test_case.index;
  gtl::InlinedVector<TensorValue, 4> inputs({TensorValue(&range_dataset_tensor),
                                             TensorValue(&num_shards),
                                             TensorValue(&index)});
  std::unique_ptr<OpKernelContext> shard_dataset_context;
  TF_ASSERT_OK(CreateShardDatasetContext(shard_dataset_kernel.get(), &inputs,
                                         &shard_dataset_context));

  DatasetBase* shard_dataset;
  TF_ASSERT_OK(CreateDataset(shard_dataset_kernel.get(),
                             shard_dataset_context.get(), &shard_dataset));
  core::ScopedUnref scoped_unref_batch_dataset(shard_dataset);

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

  bool end_of_sequence = false;
  auto expected_outputs_it = test_case.expected_outputs.begin();
  std::vector<Tensor> out_tensors;
  while (!end_of_sequence) {
    TF_EXPECT_OK(
        iterator->GetNext(iterator_ctx.get(), &out_tensors, &end_of_sequence));
    if (!end_of_sequence) {
      EXPECT_LT(expected_outputs_it, test_case.expected_outputs.end());
      TF_EXPECT_OK(ExpectEqual(out_tensors.back(), *expected_outputs_it));
      expected_outputs_it++;
    }
  }
  EXPECT_EQ(expected_outputs_it, test_case.expected_outputs.end());
}

TEST_F(ShardDatasetOpTest, DatasetNodeName) {
  int thread_num = 2, cpu_num = 2;