| 309 | public ::testing::WithParamInterface<TestCase> {}; |
| 310 | |
| 311 | TEST_P(ParameterizedShuffleDatasetOpTest, GetNext) { |
| 312 | int thread_num = 2, cpu_num = 2; |
| 313 | TestCase test_case = GetParam(); |
| 314 | TF_ASSERT_OK(InitThreadPool(thread_num)); |
| 315 | TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); |
| 316 | |
| 317 | Tensor count = test_case.count; |
| 318 | int64 count_value = count.flat<int64>()(0); |
| 319 | std::unique_ptr<OpKernel> dataset_kernel; |
| 320 | TF_ASSERT_OK( |
| 321 | CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration, |
| 322 | test_case.expected_output_dtypes, |
| 323 | test_case.expected_output_shapes, &dataset_kernel)); |
| 324 | |
| 325 | DatasetBase* range_dataset; |
| 326 | TF_ASSERT_OK(CreateRangeDataset<int64>( |
| 327 | test_case.range_data_param.start, test_case.range_data_param.end, |
| 328 | test_case.range_data_param.step, "range", &range_dataset)); |
| 329 | Tensor range_dataset_tensor(DT_VARIANT, TensorShape({})); |
| 330 | TF_ASSERT_OK( |
| 331 | StoreDatasetInVariantTensor(range_dataset, &range_dataset_tensor)); |
| 332 | Tensor buffer_size = test_case.buffer_size; |
| 333 | Tensor seed = test_case.seed; |
| 334 | Tensor seed2 = test_case.seed2; |
| 335 | gtl::InlinedVector<TensorValue, 4> inputs( |
| 336 | {TensorValue(&range_dataset_tensor), TensorValue(&buffer_size), |
| 337 | TensorValue(&seed), TensorValue(&seed2)}); |
| 338 | if (count_value != 1) inputs.push_back(TensorValue(&count)); |
| 339 | |
| 340 | std::unique_ptr<OpKernelContext> dataset_context; |
| 341 | TF_ASSERT_OK( |
| 342 | CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context)); |
| 343 | DatasetBase* dataset; |
| 344 | TF_ASSERT_OK( |
| 345 | CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset)); |
| 346 | core::ScopedUnref scoped_unref_dataset(dataset); |
| 347 | |
| 348 | std::unique_ptr<IteratorContext> iterator_ctx; |
| 349 | TF_ASSERT_OK(CreateIteratorContext(dataset_context.get(), &iterator_ctx)); |
| 350 | std::unique_ptr<IteratorBase> iterator; |
| 351 | TF_ASSERT_OK( |
| 352 | dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator)); |
| 353 | |
| 354 | bool end_of_sequence = false; |
| 355 | std::vector<Tensor> shuffled_out_tensors; |
| 356 | while (!end_of_sequence) { |
| 357 | std::vector<Tensor> next; |
| 358 | TF_EXPECT_OK( |
| 359 | iterator->GetNext(iterator_ctx.get(), &next, &end_of_sequence)); |
| 360 | shuffled_out_tensors.insert(shuffled_out_tensors.end(), next.begin(), |
| 361 | next.end()); |
| 362 | // For the forever-repeat case, we test only a finite number of steps of |
| 363 | // the infinite sequence. |
| 364 | if (count_value == -1 && shuffled_out_tensors.size() == |
| 365 | test_case.expected_shuffle_outputs.size()) { |
| 366 | break; |
| 367 | } |
| 368 | } |
nothing calls this directly
no test coverage detected