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

     1  import unittest
     2  from unittest.mock import Mock
     3  
     4  from torch.utils.data import IterDataPipe
     5  
     6  from aistore.pytorch.aisio import AISSourceLister
     7  from aistore.sdk.ais_source import AISSource
     8  from tests.const import ETL_NAME
     9  
    10  
    11  class TestDataPipes(unittest.TestCase):
    12      def test_source_lister(self):
    13          ais_source_1 = Mock(AISSource)
    14          ais_source_2 = Mock(AISSource)
    15          source_1_urls = ["url1", "url2"]
    16          source_2_urls = ["url3", "url4"]
    17          ais_source_1.list_urls.return_value = source_1_urls
    18          ais_source_2.list_urls.return_value = source_2_urls
    19          expected_res = source_1_urls + source_2_urls
    20          prefix = "obj-prefix-"
    21  
    22          source_lister = AISSourceLister(
    23              [ais_source_1, ais_source_2], prefix=prefix, etl_name=ETL_NAME
    24          )
    25  
    26          self.assertIsInstance(source_lister, IterDataPipe)
    27          self.assertEqual(expected_res, list(source_lister))
    28          ais_source_1.list_urls.assert_called_with(prefix=prefix, etl_name=ETL_NAME)
    29          ais_source_2.list_urls.assert_called_with(prefix=prefix, etl_name=ETL_NAME)