github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/python/tests/unit/sdk/test_object_iterator.py (about)

     1  import unittest
     2  from unittest.mock import Mock
     3  
     4  from aistore.sdk.object_iterator import ObjectIterator
     5  from aistore.sdk.types import BucketEntry
     6  
     7  
     8  class TestObjectIterator(unittest.TestCase):  # pylint: disable=unused-variable
     9      def setUp(self) -> None:
    10          self.callable_resp = Mock()
    11          self.callable_resp.continuation_token = ""
    12          self.callable = lambda uuid, token: self.callable_resp
    13          self.obj_iterator = ObjectIterator(self.callable)
    14  
    15      def test_iter(self):
    16          self.assertEqual(self.obj_iterator, iter(self.obj_iterator))
    17  
    18      def test_next_empty_resp(self):
    19          with self.assertRaises(StopIteration):
    20              self.callable_resp.entries = []
    21              self.callable_resp.uuid = ""
    22              self.callable_resp.continuation_token = ""
    23              next(self.obj_iterator)
    24  
    25      def test_next_iterator_exhausted(self):
    26          entry_1 = Mock(BucketEntry)
    27          entry_2 = Mock(BucketEntry)
    28          entry_3 = Mock(BucketEntry)
    29          self.callable_resp.entries = [entry_1, entry_2, entry_3]
    30          self.callable_resp.uuid = "UUID"
    31          self.assertEqual(entry_1, next(self.obj_iterator))
    32          self.assertEqual(entry_2, next(self.obj_iterator))
    33          self.assertEqual(entry_3, next(self.obj_iterator))
    34          with self.assertRaises(StopIteration):
    35              next(self.obj_iterator)
    36  
    37      def test_next_multiple_pages(self):
    38          entry_1 = Mock(BucketEntry)
    39          entry_2 = Mock(BucketEntry)
    40          entry_3 = Mock(BucketEntry)
    41          self.callable_resp.entries = [entry_1, entry_2]
    42          self.callable_resp.uuid = ""
    43          self.assertEqual(entry_1, next(self.obj_iterator))
    44          self.assertEqual(entry_2, next(self.obj_iterator))
    45          self.callable_resp.entries = [entry_3]
    46          self.assertEqual(entry_3, next(self.obj_iterator))
    47          with self.assertRaises(StopIteration):
    48              next(self.obj_iterator)