From f4fde40524d73d855bb5ea6375834dce24cd4688 Mon Sep 17 00:00:00 2001
From: Mukund Thakur <mthakur@cloudera.com>
Date: Tue, 28 May 2024 11:27:33 -0500
Subject: [PATCH] HADOOP-19184. S3A Fix
 TestStagingCommitter.testJobCommitFailure (#6843)

Follow up on HADOOP-18679

Contributed by: Mukund Thakur
---
 .../apache/hadoop/fs/s3a/MockS3AFileSystem.java |  6 +++++-
 .../fs/s3a/commit/staging/StagingTestBase.java  | 17 +++++++++++++----
 .../commit/staging/TestStagingCommitter.java    |  2 +-
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/hadoop-tools/hadoop-aws/src/test/java/org/apache/hadoop/fs/s3a/MockS3AFileSystem.java b/hadoop-tools/hadoop-aws/src/test/java/org/apache/hadoop/fs/s3a/MockS3AFileSystem.java
index b7e55f01a3..bf3fd27701 100644
--- a/hadoop-tools/hadoop-aws/src/test/java/org/apache/hadoop/fs/s3a/MockS3AFileSystem.java
+++ b/hadoop-tools/hadoop-aws/src/test/java/org/apache/hadoop/fs/s3a/MockS3AFileSystem.java
@@ -353,7 +353,11 @@ void deleteObjectAtPath(Path f,
       String key,
       boolean isFile)
       throws SdkException, IOException {
-    deleteObject(key);
+    mock.getS3AInternals()
+            .getAmazonS3Client("test")
+            .deleteObject(getRequestFactory()
+            .newDeleteObjectRequestBuilder(key)
+            .build());
   }
 
   @Override
diff --git a/hadoop-tools/hadoop-aws/src/test/java/org/apache/hadoop/fs/s3a/commit/staging/StagingTestBase.java b/hadoop-tools/hadoop-aws/src/test/java/org/apache/hadoop/fs/s3a/commit/staging/StagingTestBase.java
index e64822d8c8..9fba584fbd 100644
--- a/hadoop-tools/hadoop-aws/src/test/java/org/apache/hadoop/fs/s3a/commit/staging/StagingTestBase.java
+++ b/hadoop-tools/hadoop-aws/src/test/java/org/apache/hadoop/fs/s3a/commit/staging/StagingTestBase.java
@@ -43,6 +43,9 @@
 import software.amazon.awssdk.services.s3.model.MultipartUpload;
 import software.amazon.awssdk.services.s3.model.UploadPartRequest;
 import software.amazon.awssdk.services.s3.model.UploadPartResponse;
+
+import org.apache.hadoop.fs.s3a.S3AInternals;
+import org.apache.hadoop.fs.s3a.S3AStore;
 import org.apache.hadoop.util.Lists;
 import org.apache.hadoop.thirdparty.com.google.common.collect.Maps;
 import org.junit.AfterClass;
@@ -129,9 +132,10 @@ protected StagingTestBase() {
    * @throws IOException IO problems.
    */
   protected static S3AFileSystem createAndBindMockFSInstance(Configuration conf,
-      Pair<StagingTestBase.ClientResults, StagingTestBase.ClientErrors> outcome)
+      Pair<StagingTestBase.ClientResults, StagingTestBase.ClientErrors> outcome,
+                                                             S3Client mockS3Client)
       throws IOException {
-    S3AFileSystem mockFs = mockS3AFileSystemRobustly();
+    S3AFileSystem mockFs = mockS3AFileSystemRobustly(mockS3Client);
     MockS3AFileSystem wrapperFS = new MockS3AFileSystem(mockFs, outcome);
     URI uri = RAW_BUCKET_URI;
     wrapperFS.initialize(uri, conf);
@@ -142,8 +146,13 @@ protected static S3AFileSystem createAndBindMockFSInstance(Configuration conf,
     return mockFs;
   }
 
-  private static S3AFileSystem mockS3AFileSystemRobustly() {
+  private static S3AFileSystem mockS3AFileSystemRobustly(S3Client mockS3Client) {
     S3AFileSystem mockFS = mock(S3AFileSystem.class);
+    S3AInternals s3AInternals = mock(S3AInternals.class);
+    when(mockFS.getS3AInternals()).thenReturn(s3AInternals);
+    when(s3AInternals.getStore()).thenReturn(mock(S3AStore.class));
+    when(s3AInternals.getAmazonS3Client(anyString()))
+        .thenReturn(mockS3Client);
     doNothing().when(mockFS).incrementReadOperations();
     doNothing().when(mockFS).incrementWriteOperations();
     doNothing().when(mockFS).incrementWriteOperations();
@@ -350,7 +359,7 @@ public void setupJob() throws Exception {
       this.errors = new StagingTestBase.ClientErrors();
       this.mockClient = newMockS3Client(results, errors);
       this.mockFS = createAndBindMockFSInstance(jobConf,
-          Pair.of(results, errors));
+          Pair.of(results, errors), mockClient);
       this.wrapperFS = lookupWrapperFS(jobConf);
       // and bind the FS
       wrapperFS.setAmazonS3Client(mockClient);
diff --git a/hadoop-tools/hadoop-aws/src/test/java/org/apache/hadoop/fs/s3a/commit/staging/TestStagingCommitter.java b/hadoop-tools/hadoop-aws/src/test/java/org/apache/hadoop/fs/s3a/commit/staging/TestStagingCommitter.java
index 71ed0b6891..fae7a6232d 100644
--- a/hadoop-tools/hadoop-aws/src/test/java/org/apache/hadoop/fs/s3a/commit/staging/TestStagingCommitter.java
+++ b/hadoop-tools/hadoop-aws/src/test/java/org/apache/hadoop/fs/s3a/commit/staging/TestStagingCommitter.java
@@ -158,7 +158,7 @@ public void setupCommitter() throws Exception {
     this.errors = new StagingTestBase.ClientErrors();
     this.mockClient = newMockS3Client(results, errors);
     this.mockFS = createAndBindMockFSInstance(jobConf,
-        Pair.of(results, errors));
+        Pair.of(results, errors), mockClient);
     this.wrapperFS = lookupWrapperFS(jobConf);
     // and bind the FS
     wrapperFS.setAmazonS3Client(mockClient);