Function TEST

tensorflow/core/kernels/batching_util/batch_scheduler_test.cc:40–81

};

TEST(BatchTest, Basic) {
  Batch<FakeTask> batch;

  EXPECT_EQ(0, batch.num_tasks());
  EXPECT_TRUE(batch.empty());
  EXPECT_EQ(0, batch.size());
  EXPECT_FALSE(batch.IsClosed());

  auto task0 = new FakeTask(3);
  batch.AddTask(std::unique_ptr<FakeTask>(task0));

  EXPECT_EQ(1, batch.num_tasks());
  EXPECT_FALSE(batch.empty());
  EXPECT_EQ(task0->size(), batch.size());
  EXPECT_EQ(task0->size(), batch.task(0).size());
  EXPECT_FALSE(batch.IsClosed());

  auto task1 = new FakeTask(7);
  batch.AddTask(std::unique_ptr<FakeTask>(task1));

  EXPECT_EQ(2, batch.num_tasks());
  EXPECT_FALSE(batch.empty());
  EXPECT_EQ(task0->size() + task1->size(), batch.size());
  EXPECT_EQ(task1->size(), batch.task(1).size());
  EXPECT_EQ(task1->size(), batch.mutable_task(1)->size());
  EXPECT_FALSE(batch.IsClosed());

  batch.Close();
  EXPECT_TRUE(batch.IsClosed());

  EXPECT_EQ(2, batch.num_tasks());
  EXPECT_FALSE(batch.empty());
  EXPECT_EQ(task0->size() + task1->size(), batch.size());
  EXPECT_EQ(task0->size(), batch.task(0).size());
  EXPECT_EQ(task1->size(), batch.task(1).size());

  EXPECT_EQ(7, batch.RemoveTask()->size());
  EXPECT_EQ(3, batch.size());
  EXPECT_EQ(3, batch.RemoveTask()->size());
  EXPECT_EQ(0, batch.size());
  EXPECT_TRUE(batch.empty());
}

TEST(BatchTest, WaitUntilClosed) {
  Batch<FakeTask> batch;