| 125 | public ::testing::WithParamInterface<TestCase> {}; |
| 126 | |
| 127 | TEST_P(ParameterizedDatasetOpTest, GetNext) { |
| 128 | int thread_num = 2, cpu_num = 2; |
| 129 | TF_ASSERT_OK(InitThreadPool(thread_num)); |
| 130 | TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); |
| 131 | const TestCase &test_case = GetParam(); |
| 132 | Tensor tensor_slice_dataset_tensor(DT_VARIANT, TensorShape({})); |
| 133 | std::vector<Tensor> inputs_for_tensor_slice_dataset = test_case.input_tensors; |
| 134 | TF_ASSERT_OK(CreateTensorSliceDatasetTensor(&inputs_for_tensor_slice_dataset, |
| 135 | &tensor_slice_dataset_tensor)); |
| 136 | Tensor count = CreateTensor<int64>(TensorShape{}, {test_case.count}); |
| 137 | gtl::InlinedVector<TensorValue, 4> inputs_for_repeat_dataset; |
| 138 | inputs_for_repeat_dataset.emplace_back(&tensor_slice_dataset_tensor); |
| 139 | inputs_for_repeat_dataset.emplace_back(&count); |
| 140 | |
| 141 | std::unique_ptr<OpKernel> repeat_dataset_kernel; |
| 142 | TF_ASSERT_OK(CreateRepeatDatasetKernel(test_case.expected_output_dtypes, |
| 143 | test_case.expected_output_shapes, |
| 144 | &repeat_dataset_kernel)); |
| 145 | std::unique_ptr<OpKernelContext> repeat_dataset_context; |
| 146 | TF_ASSERT_OK(CreateRepeatDatasetContext(repeat_dataset_kernel.get(), |
| 147 | &inputs_for_repeat_dataset, |
| 148 | &repeat_dataset_context)); |
| 149 | DatasetBase *repeat_dataset; |
| 150 | TF_ASSERT_OK(CreateDataset(repeat_dataset_kernel.get(), |
| 151 | repeat_dataset_context.get(), &repeat_dataset)); |
| 152 | core::ScopedUnref scoped_unref(repeat_dataset); |
| 153 | |
| 154 | std::unique_ptr<IteratorContext> iterator_ctx; |
| 155 | TF_ASSERT_OK( |
| 156 | CreateIteratorContext(repeat_dataset_context.get(), &iterator_ctx)); |
| 157 | std::unique_ptr<IteratorBase> iterator; |
| 158 | TF_ASSERT_OK( |
| 159 | repeat_dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator)); |
| 160 | |
| 161 | auto expected_outputs_it = test_case.expected_outputs.begin(); |
| 162 | bool end_of_sequence = false; |
| 163 | std::vector<Tensor> out_tensors; |
| 164 | |
| 165 | if (test_case.count < 0) { |
| 166 | // We test only a finite number of steps of the infinite sequence. |
| 167 | for (int i = 0; i < 100; ++i) { |
| 168 | out_tensors.clear(); |
| 169 | TF_EXPECT_OK(iterator->GetNext(iterator_ctx.get(), &out_tensors, |
| 170 | &end_of_sequence)); |
| 171 | for (const auto &tensor : out_tensors) { |
| 172 | TF_EXPECT_OK(ExpectEqual(tensor, *expected_outputs_it)); |
| 173 | expected_outputs_it++; |
| 174 | // In the forever-repeat test case, the first group of the repeated |
| 175 | // tensors is used to represent the expected outputs, so the iterator |
| 176 | // of the expected outputs needs to be reset once it reaches the end. |
| 177 | if (expected_outputs_it == test_case.expected_outputs.end()) { |
| 178 | expected_outputs_it = test_case.expected_outputs.begin(); |
| 179 | } |
| 180 | } |
| 181 | } |
| 182 | EXPECT_FALSE(end_of_sequence); |
| 183 | } else { |
| 184 | while (!end_of_sequence) { |
nothing calls this directly
no test coverage detected