westonpace commented on a change in pull request #9808:
URL: https://github.com/apache/arrow/pull/9808#discussion_r605143617



##########
File path: cpp/src/arrow/util/async_generator_test.cc
##########
@@ -570,14 +659,134 @@ TEST(TestAsyncUtil, StackOverflow) {
 
 #endif
 
-TEST(TestAsyncUtil, Background) {
+class BackgroundGeneratorTestFixture : public GeneratorTestFixture {
+ protected:
+  AsyncGenerator<TestInt> Make(const std::vector<TestInt>& it,
+                               int max_q = kDefaultBackgroundMaxQ,
+                               int q_restart = kDefaultBackgroundQRestart) {
+    bool slow = GetParam();
+    return BackgroundAsyncVectorIt(it, slow, max_q, q_restart);
+  }
+};
+
+TEST_P(BackgroundGeneratorTestFixture, Empty) {
+  auto background = Make({});
+  AssertGeneratorExhausted(background);
+}
+
+TEST_P(BackgroundGeneratorTestFixture, Basic) {
   std::vector<TestInt> expected = {1, 2, 3};
-  auto background = BackgroundAsyncVectorIt(expected);
+  auto background = Make(expected);
   auto future = CollectAsyncGenerator(background);
   ASSERT_FINISHES_OK_AND_ASSIGN(auto collected, future);
   ASSERT_EQ(expected, collected);
 }
 
+TEST_P(BackgroundGeneratorTestFixture, BadResult) {
+  std::shared_ptr<ManualIteratorControl<TestInt>> iterator_control;
+  auto iterator = MakePushIterator<TestInt>(&iterator_control);
+  // Enough valid items to fill the queue and then some
+  for (int i = 0; i < 5; i++) {
+    iterator_control->Push(i);
+  }
+  // Next fail
+  iterator_control->Push(Status::Invalid("XYZ"));
+  ASSERT_OK_AND_ASSIGN(
+      auto generator,
+      MakeBackgroundGenerator(std::move(iterator), 
internal::GetCpuThreadPool(), 4, 2));
+
+  ASSERT_FINISHES_OK_AND_EQ(TestInt(0), generator());
+  // Have not yet restarted so next results should always be valid
+  ASSERT_FINISHES_OK_AND_EQ(TestInt(1), generator());
+  ASSERT_FINISHES_OK_AND_EQ(TestInt(2), generator());

Review comment:
       If it restarts it will restart on the call for 2 and the call that 
restarts can never fail in this test (because it doesn't restart until after 
it's grabbed a result from the queue).
   
   I think what is happening is things are slowing down in just the right way 
so the background reader never actually fills the queue and so 5 is able to 
sneak in before 2 is asked for.
   
   Either way, the fix is the same, I now allow the call for 2 to fail.  Thanks 
for catching this.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
[email protected]


Reply via email to