Function TEST_P

tensorflow/core/kernels/data/tensor_dataset_op_test.cc:122–170

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

TEST_P(ParametrizedTensorDatasetOpTest, 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();
  std::vector<Tensor> components = test_case.components;
  gtl::InlinedVector<TensorValue, 4> inputs;
  for (auto &component : components) {
    inputs.push_back(TensorValue(&component));
  }
  std::unique_ptr<OpKernel> tensor_dataset_kernel;
  TF_ASSERT_OK(CreateTensorDatasetKernel(test_case.expected_output_dtypes,
                                         test_case.expected_output_shapes,
                                         &tensor_dataset_kernel));
  std::unique_ptr<OpKernelContext> tensor_dataset_context;
  TF_ASSERT_OK(CreateTensorDatasetContext(tensor_dataset_kernel.get(), &inputs,
                                          &tensor_dataset_context));
  DatasetBase *tensor_dataset;
  TF_ASSERT_OK(CreateDataset(tensor_dataset_kernel.get(),
                             tensor_dataset_context.get(), &tensor_dataset));
  core::ScopedUnref scoped_unref(tensor_dataset);

  std::unique_ptr<IteratorContext> iterator_context;
  TF_ASSERT_OK(
      CreateIteratorContext(tensor_dataset_context.get(), &iterator_context));
  std::unique_ptr<IteratorBase> iterator;
  TF_ASSERT_OK(tensor_dataset->MakeIterator(iterator_context.get(), "Iterator",
                                            &iterator));
  bool end_of_sequence = false;
  std::vector<Tensor> out_tensors;
  while (!end_of_sequence) {
    TF_EXPECT_OK(iterator->GetNext(iterator_context.get(), &out_tensors,
                                   &end_of_sequence));
  }
  EXPECT_EQ(out_tensors.size(), test_case.expected_outputs.size());
  for (int i = 0; i < out_tensors.size(); ++i) {
    if (out_tensors[i].dtype() == DT_VARIANT) {
      // Currently `ExpectEqual()` does not support the variant tensor
      // yet, so we manually cast the variant to numeric/string tensor.
      const Tensor *output = out_tensors[i].scalar<Variant>()().get<Tensor>();
      const Tensor *expected_output =
          test_case.expected_outputs[i].scalar<Variant>()().get<Tensor>();
      TF_EXPECT_OK(ExpectEqual(*output, *expected_output));
    } else {
      TF_EXPECT_OK(ExpectEqual(out_tensors[i], test_case.expected_outputs[i]));
    }
  }
}

TEST_F(TensorDatasetOpTest, DatasetTypeString) {
  int thread_num = 2, cpu_num = 2;