Function TEST_P

tensorflow/core/kernels/data/take_dataset_op_test.cc:150–199

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

TEST_P(ParameterizedTakeDatasetOpTest, GetNext) {
  int thread_num = 2, cpu_num = 2;
  TF_ASSERT_OK(InitThreadPool(thread_num));
  TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));
  const TestCase &test_case = GetParam();
  Tensor tensor_slice_dataset_tensor(DT_VARIANT, TensorShape({}));
  std::vector<Tensor> inputs_for_tensor_slice_dataset = test_case.input_tensors;
  TF_ASSERT_OK(CreateTensorSliceDatasetTensor(&inputs_for_tensor_slice_dataset,
                                              &tensor_slice_dataset_tensor));
  Tensor count = CreateTensor<int64>(TensorShape{}, {test_case.count});
  gtl::InlinedVector<TensorValue, 4> inputs_for_take_dataset;
  inputs_for_take_dataset.emplace_back(&tensor_slice_dataset_tensor);
  inputs_for_take_dataset.emplace_back(&count);

  std::unique_ptr<OpKernel> take_dataset_kernel;
  TF_ASSERT_OK(CreateTakeDatasetKernel(test_case.expected_output_dtypes,
                                       test_case.expected_output_shapes,
                                       &take_dataset_kernel));
  std::unique_ptr<OpKernelContext> take_dataset_context;
  TF_ASSERT_OK(CreateTakeDatasetContext(take_dataset_kernel.get(),
                                        &inputs_for_take_dataset,
                                        &take_dataset_context));
  DatasetBase *take_dataset;
  TF_ASSERT_OK(CreateDataset(take_dataset_kernel.get(),
                             take_dataset_context.get(), &take_dataset));
  core::ScopedUnref scoped_unref(take_dataset);

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

  auto expected_outputs_it = test_case.expected_outputs.begin();
  bool end_of_sequence = false;
  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) {
      for (const auto &tensor : out_tensors) {
        EXPECT_NE(expected_outputs_it, test_case.expected_outputs.end());
        TF_EXPECT_OK(ExpectEqual(tensor, *expected_outputs_it));
        expected_outputs_it++;
      }
    }
  }
  EXPECT_EQ(expected_outputs_it, test_case.expected_outputs.end());
}

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