def test_batching_no_remainder(self):
    """Tests if batching does not drop any example."""
    index_size, num_examples, batch_size = 32, 24, 8
    index = [{"hashed_id": str(i)} for i in range(index_size)]
    retriever = _MockEmbeddingBasedRetriever(index, batch_size=batch_size)
    examples = [{"hashed_id": str(i)} for i in range(num_examples)]
    results = list(retriever.retrieve_all(examples))
    self.assertLen(results, num_examples)
    for i, result in enumerate(results):
      self.assertEqual(result.example["hashed_id"], str(i))
      self.assertLen(result.neighbor_ids, index_size)
      self.assertLen(result.neighbor_distances, index_size)
      for j in range(1, index_size):
        # Higher-ranked neighbors should have smaller distances.
        self.assertGreaterEqual(result.neighbor_distances[j],
                                result.neighbor_distances[j - 1])
Tests if batching does not drop any example.
language/casper/retrieve/query_retrievers_test.py