Function TEST_P

tensorflow/core/kernels/data/cache_dataset_ops_test.cc:154–212

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

TEST_P(ParameterizedCacheDatasetOpTest, 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> cache_dataset_kernel;
  TF_ASSERT_OK(CreateCacheDatasetOpKernel(test_case.expected_output_dtypes,
                                          test_case.expected_output_shapes,
                                          &cache_dataset_kernel));
  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 file_name =
      CreateTensor<tstring>(TensorShape{}, {test_case.file_name});
  gtl::InlinedVector<TensorValue, 4> inputs(
      {TensorValue(&tensor_slice_dataset_tensor), TensorValue(&file_name)});
  std::unique_ptr<OpKernelContext> cache_dataset_context;
  TF_ASSERT_OK(CreateCacheDatasetContext(cache_dataset_kernel.get(), &inputs,
                                         &cache_dataset_context));
  DatasetBase* cache_dataset;
  TF_ASSERT_OK(CreateDataset(cache_dataset_kernel.get(),
                             cache_dataset_context.get(), &cache_dataset));
  core::ScopedUnref scoped_unref(cache_dataset);

  std::unique_ptr<IteratorContext> iterator_ctx;
  TF_ASSERT_OK(
      CreateIteratorContext(cache_dataset_context.get(), &iterator_ctx));
  std::unique_ptr<IteratorBase> iterator;
  TF_ASSERT_OK(cache_dataset->MakeIterator(iterator_ctx.get(), kIteratorPrefix,
                                           &iterator));

  // Test the write mode.
  bool end_of_sequence = false;
  std::vector<Tensor> out_tensors;
  while (!end_of_sequence) {
    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());
  }
  TF_EXPECT_OK(ExpectEqual(out_tensors, test_case.expected_outputs,
                           /*compare_order*/ true));

  // Test the read mode.
  TF_ASSERT_OK(cache_dataset->MakeIterator(iterator_ctx.get(), kIteratorPrefix,
                                           &iterator));
  end_of_sequence = false;
  out_tensors.clear();
  while (!end_of_sequence) {
    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());
  }
  TF_EXPECT_OK(ExpectEqual(out_tensors, test_case.expected_outputs,
                           /*compare_order*/ true));