From 379cd3bf14cce84c1be45eb7f9dcd8cdb2e3f9c9 Mon Sep 17 00:00:00 2001
From: Roman Chernykh <r.chernykh@yadro.com>
Date: Mon, 18 Nov 2024 13:01:26 +0300
Subject: [PATCH] [#324]Extension list_objects method

---
 src/frostfs_testlib/s3/aws_cli_client.py | 13 +++++++++++--
 src/frostfs_testlib/s3/boto3_client.py   | 11 +++++++++--
 src/frostfs_testlib/s3/interfaces.py     |  4 +++-
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/src/frostfs_testlib/s3/aws_cli_client.py b/src/frostfs_testlib/s3/aws_cli_client.py
index ba95733..2ac6d68 100644
--- a/src/frostfs_testlib/s3/aws_cli_client.py
+++ b/src/frostfs_testlib/s3/aws_cli_client.py
@@ -196,11 +196,20 @@ class AwsCliClient(S3ClientWrapper):
         return response.get("LocationConstraint")
 
     @reporter.step("List objects S3")
-    def list_objects(self, bucket: str, full_output: bool = False) -> Union[dict, list[str]]:
+    def list_objects(
+        self, bucket: str, full_output: bool = False, page_size: Optional[int] = None, prefix: Optional[str] = None
+    ) -> Union[dict, list[str]]:
         if bucket.startswith("-") or " " in bucket:
             bucket = f'"{bucket}"'
 
-        cmd = f"aws {self.common_flags} s3api list-objects --bucket {bucket} " f"--endpoint {self.s3gate_endpoint} --profile {self.profile}"
+        cmd = f"aws {self.common_flags} s3api list-objects --bucket {bucket} " f"--endpoint {self.s3gate_endpoint} "
+        if page_size:
+            cmd = cmd.replace("--no-paginate", "")
+            cmd += f" --page-size {page_size} "
+        if prefix:
+            cmd += f" --prefix {prefix}"
+        if self.profile:
+            cmd += f" --profile {self.profile} "
         output = self.local_shell.exec(cmd).stdout
         response = self._to_json(output)
 
diff --git a/src/frostfs_testlib/s3/boto3_client.py b/src/frostfs_testlib/s3/boto3_client.py
index 12113ad..e7f2c35 100644
--- a/src/frostfs_testlib/s3/boto3_client.py
+++ b/src/frostfs_testlib/s3/boto3_client.py
@@ -398,10 +398,17 @@ class Boto3ClientWrapper(S3ClientWrapper):
         return response if full_output else obj_list
 
     @reporter.step("List objects S3")
-    def list_objects(self, bucket: str, full_output: bool = False) -> Union[dict, list[str]]:
+    def list_objects(
+        self, bucket: str, full_output: bool = False, page_size: Optional[int] = None, prefix: Optional[str] = None
+    ) -> Union[dict, list[str]]:
+        params = {"Bucket": bucket}
+        if page_size:
+            params["MaxKeys"] = page_size
+        if prefix:
+            params["Prefix"] = prefix
         response = self._exec_request(
             self.boto3_client.list_objects,
-            params={"Bucket": bucket},
+            params,
             endpoint=self.s3gate_endpoint,
             profile=self.profile,
         )
diff --git a/src/frostfs_testlib/s3/interfaces.py b/src/frostfs_testlib/s3/interfaces.py
index 69a5154..c3d99eb 100644
--- a/src/frostfs_testlib/s3/interfaces.py
+++ b/src/frostfs_testlib/s3/interfaces.py
@@ -195,7 +195,9 @@ class S3ClientWrapper(HumanReadableABC):
         """
 
     @abstractmethod
-    def list_objects(self, bucket: str, full_output: bool = False) -> Union[dict, list[str]]:
+    def list_objects(
+        self, bucket: str, full_output: bool = False, page_size: Optional[int] = None, prefix: Optional[str] = None
+    ) -> Union[dict, list[str]]:
         """Returns some or all (up to 1,000) of the objects in a bucket with each request.
         You can use the request parameters as selection criteria to return a subset of the objects in a bucket.
         A 200 OK response can contain valid or invalid XML. Make sure to design your application