From 5d4b0897151bdd81571960b7a5e7acebc680fce9 Mon Sep 17 00:00:00 2001
From: mosiac <mosiac@bloomberg.net>
Date: Mon, 23 Sep 2024 16:57:39 +0100
Subject: [PATCH] Add test for exception propagation from plugins

---
 ...edentialsProviderExceptionPropagation.java | 82 +++++++++++++++++++
 1 file changed, 82 insertions(+)
 create mode 100644 trino-aws-proxy/src/test/java/io/trino/aws/proxy/server/plugin/exception/TestCredentialsProviderExceptionPropagation.java

diff --git a/trino-aws-proxy/src/test/java/io/trino/aws/proxy/server/plugin/exception/TestCredentialsProviderExceptionPropagation.java b/trino-aws-proxy/src/test/java/io/trino/aws/proxy/server/plugin/exception/TestCredentialsProviderExceptionPropagation.java
new file mode 100644
index 00000000..c7f4fc91
--- /dev/null
+++ b/trino-aws-proxy/src/test/java/io/trino/aws/proxy/server/plugin/exception/TestCredentialsProviderExceptionPropagation.java
@@ -0,0 +1,82 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.aws.proxy.server.plugin.exception;
+
+import com.google.inject.Inject;
+import com.google.inject.Scopes;
+import io.trino.aws.proxy.server.testing.TestingTrinoAwsProxyServer.Builder;
+import io.trino.aws.proxy.server.testing.harness.BuilderFilter;
+import io.trino.aws.proxy.server.testing.harness.TrinoAwsProxyTest;
+import io.trino.aws.proxy.spi.credentials.Credentials;
+import io.trino.aws.proxy.spi.credentials.CredentialsProvider;
+import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.services.s3.S3Client;
+import software.amazon.awssdk.services.s3.model.S3Exception;
+
+import java.util.Optional;
+
+import static io.trino.aws.proxy.spi.plugin.TrinoAwsProxyServerBinding.credentialsProviderModule;
+import static java.util.Objects.requireNonNull;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.assertj.core.api.InstanceOfAssertFactories.type;
+
+@TrinoAwsProxyTest(filters = TestCredentialsProviderExceptionPropagation.Filter.class)
+public class TestCredentialsProviderExceptionPropagation
+{
+    private final DelegatingCredentialsProvider delegatingCredentialsProvider;
+    private final S3Client internalClient;
+
+    public static class DelegatingCredentialsProvider
+            implements CredentialsProvider
+    {
+        private CredentialsProvider delegate;
+
+        public void setDelegate(CredentialsProvider delegate)
+        {
+            this.delegate = delegate;
+        }
+
+        @Override
+        public Optional<Credentials> credentials(String emulatedAccessKey, Optional<String> session)
+        {
+            return requireNonNull(delegate, "delegate is null").credentials(emulatedAccessKey, session);
+        }
+    }
+
+    public static class Filter
+            implements BuilderFilter
+    {
+        @Override
+        public Builder filter(Builder builder)
+        {
+            return builder.withoutTestingCredentialsRoleProviders()
+                    .addModule(credentialsProviderModule("testing", DelegatingCredentialsProvider.class, binder -> binder.bind(DelegatingCredentialsProvider.class).in(Scopes.SINGLETON)))
+                    .withProperty("credentials-provider.type", "testing");
+        }
+    }
+
+    @Inject
+    public TestCredentialsProviderExceptionPropagation(DelegatingCredentialsProvider delegatingCredentialsProvider, S3Client internalClient)
+    {
+        this.delegatingCredentialsProvider = requireNonNull(delegatingCredentialsProvider, "delegatingCredentialsProvider is null");
+        this.internalClient = requireNonNull(internalClient, "internalClient is null");
+    }
+
+    @Test
+    public void testExceptions()
+    {
+        delegatingCredentialsProvider.setDelegate((_, _) -> { throw new RuntimeException("Testing exception"); });
+        assertThatThrownBy(internalClient::listBuckets).asInstanceOf(type(S3Exception.class)).extracting(S3Exception::statusCode).isEqualTo(500);
+    }
+}