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

     1  import unittest
     2  from unittest.mock import Mock, patch, call
     3  
     4  from aistore.sdk import Bucket
     5  from aistore.sdk.const import (
     6      HTTP_METHOD_DELETE,
     7      ACT_DELETE_OBJECTS,
     8      ACT_EVICT_OBJECTS,
     9      HTTP_METHOD_POST,
    10      ACT_PREFETCH_OBJECTS,
    11      ACT_COPY_OBJECTS,
    12      ACT_TRANSFORM_OBJECTS,
    13      ACT_ARCHIVE_OBJECTS,
    14      PROVIDER_AMAZON,
    15      HTTP_METHOD_PUT,
    16  )
    17  from aistore.sdk.etl_const import DEFAULT_ETL_TIMEOUT
    18  from aistore.sdk.multiobj import ObjectGroup, ObjectRange
    19  from aistore.sdk.types import Namespace, BucketModel, ArchiveMultiObj
    20  from tests.const import LARGE_FILE_SIZE, ETL_NAME, PREFIX_NAME
    21  
    22  
    23  # pylint: disable=unused-variable,too-many-instance-attributes
    24  class TestObjectGroup(unittest.TestCase):
    25      def setUp(self) -> None:
    26          self.mock_bck = Mock()
    27          self.mock_bck.name = "mock-bucket"
    28          self.mock_bck.provider = "mock-bck-provider"
    29          self.mock_response_text = "Response Text"
    30          mock_response = Mock()
    31          mock_response.text = self.mock_response_text
    32          self.mock_bck.make_request.return_value = mock_response
    33          self.mock_bck_model = BucketModel(
    34              name=self.mock_bck.name, provider=self.mock_bck.provider
    35          )
    36          self.mock_bck.as_model.return_value = self.mock_bck_model
    37          namespace = Namespace(name="ns-name", uuid="ns-id")
    38          provider = "any provider"
    39          self.dest_bucket = Bucket(
    40              name="to-bucket", namespace=namespace, provider=provider
    41          )
    42  
    43          self.obj_names = ["obj-1", "obj-2"]
    44          self.object_group = ObjectGroup(self.mock_bck, obj_names=self.obj_names)
    45          self.expected_value = {}
    46          self.expected_value["objnames"] = self.obj_names
    47  
    48      def test_object_group_parameters(self):
    49          obj_names = ["list", "of", "names"]
    50          obj_range = ObjectRange(prefix=PREFIX_NAME)
    51          obj_template = "prefix-{0..3}"
    52          with self.assertRaises(ValueError):
    53              ObjectGroup(
    54                  self.mock_bck,
    55                  obj_names=obj_names,
    56                  obj_range=obj_range,
    57              )
    58          with self.assertRaises(ValueError):
    59              ObjectGroup(
    60                  self.mock_bck,
    61                  obj_names=obj_names,
    62                  obj_template=obj_template,
    63              )
    64          with self.assertRaises(ValueError):
    65              ObjectGroup(
    66                  self.mock_bck,
    67                  obj_range=obj_range,
    68                  obj_template=obj_template,
    69              )
    70  
    71      def object_group_test_helper(
    72          self, object_group_function, http_method, action, expected_value, **kwargs
    73      ):
    74          resp_text = object_group_function(**kwargs)
    75          self.assertEqual(self.mock_response_text, resp_text)
    76          self.mock_bck.make_request.assert_called_with(
    77              http_method,
    78              action,
    79              value=expected_value,
    80          )
    81  
    82      def test_delete(self):
    83          self.object_group_test_helper(
    84              self.object_group.delete,
    85              HTTP_METHOD_DELETE,
    86              ACT_DELETE_OBJECTS,
    87              self.expected_value,
    88          )
    89  
    90      def test_evict(self):
    91          self.object_group_test_helper(
    92              self.object_group.evict,
    93              HTTP_METHOD_DELETE,
    94              ACT_EVICT_OBJECTS,
    95              self.expected_value,
    96          )
    97  
    98      def test_prefetch(self):
    99          prefetch_expected_val = self.expected_value.copy()
   100          prefetch_expected_val["coer"] = False
   101          prefetch_expected_val["latest-ver"] = False
   102          self.object_group_test_helper(
   103              self.object_group.prefetch,
   104              HTTP_METHOD_POST,
   105              ACT_PREFETCH_OBJECTS,
   106              prefetch_expected_val,
   107          )
   108  
   109      def test_prefetch_with_blob_threshold(self):
   110          prefetch_expected_val = self.expected_value.copy()
   111          prefetch_expected_val["coer"] = False
   112          prefetch_expected_val["latest-ver"] = False
   113          blob_threshold_value = LARGE_FILE_SIZE
   114          prefetch_expected_val["blob-threshold"] = blob_threshold_value
   115  
   116          self.object_group_test_helper(
   117              self.object_group.prefetch,
   118              HTTP_METHOD_POST,
   119              ACT_PREFETCH_OBJECTS,
   120              prefetch_expected_val,
   121              blob_threshold=blob_threshold_value,
   122          )
   123  
   124      def test_copy(self):
   125          self.expected_value["prefix"] = ""
   126          self.expected_value["prepend"] = ""
   127          self.expected_value["dry_run"] = False
   128          self.expected_value["force"] = False
   129          self.expected_value["tobck"] = self.dest_bucket.as_model()
   130          self.expected_value["coer"] = False
   131          self.expected_value["latest-ver"] = False
   132          self.expected_value["synchronize"] = False
   133          # Test default args
   134          self.object_group_test_helper(
   135              self.object_group.copy,
   136              HTTP_METHOD_POST,
   137              ACT_COPY_OBJECTS,
   138              self.expected_value,
   139              to_bck=self.dest_bucket,
   140          )
   141          # Test provided optional args
   142          prepend_val = "new_prefix-"
   143          self.expected_value["prepend"] = prepend_val
   144          self.expected_value["force"] = True
   145          self.expected_value["dry_run"] = True
   146          self.expected_value["coer"] = True
   147          self.expected_value["latest-ver"] = False
   148          self.expected_value["synchronize"] = False
   149  
   150          self.object_group_test_helper(
   151              self.object_group.copy,
   152              HTTP_METHOD_POST,
   153              ACT_COPY_OBJECTS,
   154              self.expected_value,
   155              to_bck=self.dest_bucket,
   156              prepend=prepend_val,
   157              force=True,
   158              dry_run=True,
   159              continue_on_error=True,
   160          )
   161  
   162      @patch("aistore.sdk.multiobj.object_group.logging")
   163      def test_copy_dry_run(self, mock_logging):
   164          mock_logger = Mock()
   165          mock_logging.getLogger.return_value = mock_logger
   166  
   167          self.object_group.copy(to_bck=self.dest_bucket, dry_run=True)
   168  
   169          mock_logger.info.assert_called()
   170  
   171      def test_transform(self):
   172          self.expected_value["prefix"] = ""
   173          self.expected_value["prepend"] = ""
   174          self.expected_value["dry_run"] = False
   175          self.expected_value["force"] = False
   176          self.expected_value["id"] = ETL_NAME
   177          self.expected_value["request_timeout"] = DEFAULT_ETL_TIMEOUT
   178          self.expected_value["tobck"] = self.dest_bucket.as_model()
   179          self.expected_value["coer"] = False
   180          self.expected_value["latest-ver"] = False
   181          self.expected_value["synchronize"] = False
   182  
   183          # Test default args
   184          self.object_group_test_helper(
   185              self.object_group.transform,
   186              HTTP_METHOD_POST,
   187              ACT_TRANSFORM_OBJECTS,
   188              self.expected_value,
   189              to_bck=self.dest_bucket,
   190              etl_name=ETL_NAME,
   191          )
   192          # Test provided optional args
   193          timeout = "30s"
   194          prepend_val = "new_prefix-"
   195          self.expected_value["coer"] = True
   196          self.expected_value["prepend"] = prepend_val
   197          self.expected_value["request_timeout"] = timeout
   198          self.expected_value["dry_run"] = True
   199          self.expected_value["force"] = True
   200          self.object_group_test_helper(
   201              self.object_group.transform,
   202              HTTP_METHOD_POST,
   203              ACT_TRANSFORM_OBJECTS,
   204              self.expected_value,
   205              to_bck=self.dest_bucket,
   206              prepend=prepend_val,
   207              etl_name=ETL_NAME,
   208              timeout=timeout,
   209              dry_run=True,
   210              force=True,
   211              continue_on_error=True,
   212          )
   213  
   214      @patch("aistore.sdk.multiobj.object_group.logging")
   215      def test_transform_dry_run(self, mock_logging):
   216          mock_logger = Mock()
   217          mock_logging.getLogger.return_value = mock_logger
   218  
   219          self.object_group.transform(
   220              to_bck=self.dest_bucket, etl_name=ETL_NAME, dry_run=True
   221          )
   222  
   223          mock_logger.info.assert_called()
   224  
   225      def test_list_names(self):
   226          self.assertEqual(self.obj_names, self.object_group.list_names())
   227  
   228      def test_archive_default_params(self):
   229          archive_name = "test-arch"
   230          expected_value = ArchiveMultiObj(
   231              object_selection=self.expected_value,
   232              archive_name=archive_name,
   233              to_bck=self.mock_bck_model,
   234          ).as_dict()
   235          self.object_group_test_helper(
   236              self.object_group.archive,
   237              HTTP_METHOD_PUT,
   238              ACT_ARCHIVE_OBJECTS,
   239              expected_value=expected_value,
   240              archive_name=archive_name,
   241          )
   242  
   243      def test_archive(self):
   244          archive_name = "test-arch"
   245          namespace = Namespace(name="ns-name", uuid="ns-id")
   246          to_bck = Bucket(
   247              name="dest-bck-name", namespace=namespace, provider=PROVIDER_AMAZON
   248          )
   249          mime = "text"
   250          include_source = True
   251          allow_append = True
   252          continue_on_err = True
   253          expected_value = ArchiveMultiObj(
   254              object_selection=self.expected_value,
   255              archive_name=archive_name,
   256              to_bck=to_bck.as_model(),
   257              mime=mime,
   258              include_source_name=include_source,
   259              allow_append=allow_append,
   260              continue_on_err=continue_on_err,
   261          ).as_dict()
   262          self.object_group_test_helper(
   263              self.object_group.archive,
   264              HTTP_METHOD_PUT,
   265              ACT_ARCHIVE_OBJECTS,
   266              expected_value=expected_value,
   267              archive_name=archive_name,
   268              to_bck=to_bck,
   269              mime=mime,
   270              include_source_name=include_source,
   271              allow_append=allow_append,
   272              continue_on_err=continue_on_err,
   273          )
   274  
   275      def test_list_urls(self):
   276          expected_obj_calls = []
   277          # Should create an object reference and get url for every object returned by listing
   278          for name in self.obj_names:
   279              expected_obj_calls.append(call(name))
   280              expected_obj_calls.append(call().get_url(etl_name=ETL_NAME))
   281          list(self.object_group.list_urls(etl_name=ETL_NAME))
   282          self.mock_bck.object.assert_has_calls(expected_obj_calls)
   283  
   284      def test_list_all_objects_iter(self):
   285          res = self.object_group.list_all_objects_iter()
   286          self.assertEqual(len(list(res)), len(self.obj_names))