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

     1  import os
     2  import random
     3  import shutil
     4  import string
     5  import tempfile
     6  from pathlib import Path
     7  
     8  from aistore.sdk import Client
     9  from aistore.sdk.const import UTF_ENCODING
    10  from aistore.sdk.errors import ErrBckNotFound
    11  
    12  
    13  # pylint: disable=unused-variable
    14  def random_string(length: int = 10):
    15      return "".join(random.choices(string.ascii_lowercase, k=length))
    16  
    17  
    18  # pylint: disable=unused-variable
    19  def create_and_put_object(
    20      client: Client,
    21      bck_name: str,
    22      obj_name: str,
    23      provider: str = "ais",
    24      obj_size: int = 0,
    25  ):
    26      obj_size = obj_size if obj_size else random.randrange(10, 20)
    27      obj_body = "".join(random.choices(string.ascii_letters, k=obj_size))
    28      content = obj_body.encode(UTF_ENCODING)
    29      temp_file = Path(tempfile.gettempdir()).joinpath(os.urandom(24).hex())
    30      with open(temp_file, "wb") as file:
    31          file.write(content)
    32          file.flush()
    33          client.bucket(bck_name, provider=provider).object(obj_name).put_file(file.name)
    34      return content
    35  
    36  
    37  def destroy_bucket(client: Client, bck_name: str):
    38      try:
    39          client.bucket(bck_name).delete()
    40      except ErrBckNotFound:
    41          pass
    42  
    43  
    44  def cleanup_local(path: str):
    45      try:
    46          shutil.rmtree(path)
    47      except FileNotFoundError:
    48          pass
    49  
    50  
    51  # pylint: disable=too-many-arguments
    52  def create_and_put_objects(
    53      client, bucket, prefix, suffix, num_obj, obj_names, obj_size=None
    54  ):
    55      if not obj_names:
    56          obj_names = [prefix + str(i) + suffix for i in range(num_obj)]
    57      for obj_name in obj_names:
    58          create_and_put_object(
    59              client,
    60              bck_name=bucket.name,
    61              provider=bucket.provider,
    62              obj_name=obj_name,
    63              obj_size=obj_size,
    64          )
    65      return obj_names
    66  
    67  
    68  def test_cases(*args):
    69      def decorator(func):
    70          def wrapper(self, *inner_args, **kwargs):
    71              for arg in args:
    72                  with self.subTest(arg=arg):
    73                      func(self, arg, *inner_args, **kwargs)
    74  
    75          return wrapper
    76  
    77      return decorator