Add tests to GmailClient (#2097)

Also make GmailClient do retries on transit errors.

core/src/main/java/google/registry/groups/GmailClient.java

@@ -16,17 +16,21 @@ package google.registry.groups;
 
 import static com.google.common.collect.Iterables.toArray;
 
+import com.google.api.client.http.HttpResponseException;
 import com.google.api.services.gmail.Gmail;
 import com.google.api.services.gmail.model.Message;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.net.MediaType;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.util.EmailMessage;
 import google.registry.util.EmailMessage.Attachment;
+import google.registry.util.Retrier;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.UnsupportedEncodingException;
 import java.util.Properties;
+import java.util.function.Predicate;
 import javax.inject.Inject;
 import javax.mail.Address;
 import javax.mail.BodyPart;
@@ -38,23 +42,25 @@ import javax.mail.internet.InternetAddress;
 import javax.mail.internet.MimeBodyPart;
 import javax.mail.internet.MimeMessage;
 import javax.mail.internet.MimeMultipart;
-import org.apache.commons.codec.binary.Base64;
 
 /** Sends {@link EmailMessage EmailMessages} through Google Workspace using {@link Gmail}. */
 public final class GmailClient {
 
   private final Gmail gmail;
+  private final Retrier retrier;
   private final InternetAddress outgoingEmailAddressWithUsername;
   private final InternetAddress replyToEmailAddress;
 
   @Inject
   GmailClient(
       Gmail gmail,
+      Retrier retrier,
       @Config("gSuiteNewOutgoingEmailAddress") String gSuiteOutgoingEmailAddress,
       @Config("gSuiteOutgoingEmailDisplayName") String gSuiteOutgoingEmailDisplayName,
       @Config("replyToEmailAddress") InternetAddress replyToEmailAddress) {
 
     this.gmail = gmail;
+    this.retrier = retrier;
     this.replyToEmailAddress = replyToEmailAddress;
     try {
       this.outgoingEmailAddressWithUsername =
@@ -73,12 +79,11 @@ public final class GmailClient {
   @CanIgnoreReturnValue
   public Message sendEmail(EmailMessage emailMessage) {
     Message message = toGmailMessage(toMimeMessage(emailMessage));
-    try {
-      // "me" is reserved word for the authorized user of the Gmail API.
-      return gmail.users().messages().send("me", message).execute();
-    } catch (IOException e) {
-      throw new EmailException(e);
-    }
+    // Unlike other Cloud APIs such as GCS and SecretManager, Gmail does not retry on errors.
+    return retrier.callWithRetry(
+        // "me" is reserved word for the authorized user of the Gmail API.
+        () -> this.gmail.users().messages().send("me", message).execute(),
+        RetriableGmailExceptionPredicate.INSTANCE);
   }
 
   static Message toGmailMessage(MimeMessage message) {
@@ -86,10 +91,7 @@ public final class GmailClient {
       ByteArrayOutputStream buffer = new ByteArrayOutputStream();
       message.writeTo(buffer);
       byte[] rawMessageBytes = buffer.toByteArray();
-      String encodedEmail = Base64.encodeBase64URLSafeString(rawMessageBytes);
-      Message gmailMessage = new Message();
-      gmailMessage.setRaw(encodedEmail);
-      return gmailMessage;
+      return new Message().encodeRaw(rawMessageBytes);
     } catch (MessagingException | IOException e) {
       throw new EmailException(e);
     }
@@ -135,4 +137,41 @@ public final class GmailClient {
       super(cause);
     }
   }
+
+  /**
+   * Determines if a Gmail API exception may be retried.
+   *
+   * <p>See <a href="https://developers.google.com/gmail/api/guides/handle-errors">online doc</a>
+   * for details.
+   */
+  static class RetriableGmailExceptionPredicate implements Predicate<Throwable> {
+
+    static RetriableGmailExceptionPredicate INSTANCE = new RetriableGmailExceptionPredicate();
+
+    private static final int USAGE_LIMIT_EXCEEDED = 403;
+    private static final int TOO_MANY_REQUESTS = 429;
+    private static final int BACKEND_ERROR = 500;
+
+    private static final ImmutableSet<String> TRANSIENT_OVERAGE_REASONS =
+        ImmutableSet.of("userRateLimitExceeded", "rateLimitExceeded");
+
+    @Override
+    public boolean test(Throwable e) {
+      if (e instanceof HttpResponseException) {
+        return testHttpResponse((HttpResponseException) e);
+      }
+      return true;
+    }
+
+    private boolean testHttpResponse(HttpResponseException e) {
+      int statusCode = e.getStatusCode();
+      if (statusCode == TOO_MANY_REQUESTS || statusCode == BACKEND_ERROR) {
+        return true;
+      }
+      if (statusCode == USAGE_LIMIT_EXCEEDED) {
+        return TRANSIENT_OVERAGE_REASONS.contains(e.getStatusMessage());
+      }
+      return false;
+    }
+  }
 }

core/src/test/java/google/registry/groups/GmailClientTest.java

@@ -0,0 +1,168 @@
+// Copyright 2023 The Nomulus Authors. All Rights Reserved.
+//
+// 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 google.registry.groups;
+
+import static com.google.common.net.MediaType.CSV_UTF_8;
+import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
+import static com.google.common.truth.Truth.assertThat;
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import com.google.api.client.http.HttpResponseException;
+import com.google.api.services.gmail.Gmail;
+import com.google.api.services.gmail.model.Message;
+import google.registry.groups.GmailClient.RetriableGmailExceptionPredicate;
+import google.registry.util.EmailMessage;
+import google.registry.util.EmailMessage.Attachment;
+import google.registry.util.Retrier;
+import google.registry.util.SystemSleeper;
+import java.io.OutputStream;
+import javax.mail.Message.RecipientType;
+import javax.mail.Part;
+import javax.mail.internet.InternetAddress;
+import javax.mail.internet.MimeMessage;
+import javax.mail.internet.MimeMultipart;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.mockito.stubbing.Answer;
+import org.testcontainers.shaded.com.google.common.collect.ImmutableList;
+
+/** Unit tests for {@link GmailClient}. */
+@ExtendWith(MockitoExtension.class)
+public class GmailClientTest {
+
+  @Mock private Gmail gmail;
+  @Mock private HttpResponseException httpResponseException;
+  private GmailClient gmailClient;
+
+  @BeforeEach
+  public void setup() throws Exception {
+    gmailClient =
+        new GmailClient(
+            gmail,
+            new Retrier(new SystemSleeper(), 3),
+            "from@example.com",
+            "My sender",
+            new InternetAddress("replyTo@example.com"));
+  }
+
+  @Test
+  public void toMimeMessage_fullMessage() throws Exception {
+    InternetAddress fromAddr = new InternetAddress("from@example.com", "My sender");
+    InternetAddress toAddr = new InternetAddress("to@example.com");
+    InternetAddress ccAddr = new InternetAddress("cc@example.com");
+    InternetAddress bccAddr = new InternetAddress("bcc@example.com");
+    EmailMessage emailMessage =
+        EmailMessage.newBuilder()
+            .setFrom(fromAddr)
+            .setRecipients(ImmutableList.of(toAddr))
+            .setSubject("My subject")
+            .setBody("My body")
+            .addCc(ccAddr)
+            .addBcc(bccAddr)
+            .setAttachment(
+                Attachment.newBuilder()
+                    .setFilename("filename")
+                    .setContent("foo,bar\nbaz,qux")
+                    .setContentType(CSV_UTF_8)
+                    .build())
+            .build();
+    MimeMessage mimeMessage = gmailClient.toMimeMessage(emailMessage);
+    assertThat(mimeMessage.getFrom()).asList().containsExactly(fromAddr);
+    assertThat(mimeMessage.getRecipients(RecipientType.TO)).asList().containsExactly(toAddr);
+    assertThat(mimeMessage.getRecipients(RecipientType.CC)).asList().containsExactly(ccAddr);
+    assertThat(mimeMessage.getRecipients(RecipientType.BCC)).asList().containsExactly(bccAddr);
+    assertThat(mimeMessage.getSubject()).isEqualTo("My subject");
+    assertThat(mimeMessage.getContent()).isInstanceOf(MimeMultipart.class);
+    MimeMultipart parts = (MimeMultipart) mimeMessage.getContent();
+    Part body = parts.getBodyPart(0);
+    assertThat(body.getContentType()).isEqualTo(PLAIN_TEXT_UTF_8.toString());
+    assertThat(body.getContent()).isEqualTo("My body");
+    Part attachment = parts.getBodyPart(1);
+    assertThat(attachment.getContentType()).startsWith(CSV_UTF_8.toString());
+    assertThat(attachment.getContentType()).endsWith("name=filename");
+    assertThat(attachment.getContent()).isEqualTo("foo,bar\nbaz,qux");
+  }
+
+  @Test
+  public void toGmailMessage() throws Exception {
+    MimeMessage mimeMessage = mock(MimeMessage.class);
+    byte[] data = "My content".getBytes(UTF_8);
+    doAnswer(
+            new Answer() {
+              @Override
+              public Object answer(InvocationOnMock invocation) throws Throwable {
+                OutputStream os = invocation.getArgument(0);
+                os.write(data);
+                return null;
+              }
+            })
+        .when(mimeMessage)
+        .writeTo(any(OutputStream.class));
+    Message gmailMessage = GmailClient.toGmailMessage(mimeMessage);
+    assertThat(gmailMessage.decodeRaw()).isEqualTo(data);
+  }
+
+  @Test
+  public void isRetriable_trueIfNotHttpResponseException() {
+    assertThat(RetriableGmailExceptionPredicate.INSTANCE.test(new Exception())).isTrue();
+  }
+
+  @Test
+  public void isHttpResponseExceptionRetriable_trueIf500() {
+    when(httpResponseException.getStatusCode()).thenReturn(500);
+    assertThat(RetriableGmailExceptionPredicate.INSTANCE.test(httpResponseException)).isTrue();
+  }
+
+  @Test
+  public void isHttpResponseExceptionRetriable_trueIf429() {
+    when(httpResponseException.getStatusCode()).thenReturn(429);
+    assertThat(RetriableGmailExceptionPredicate.INSTANCE.test(httpResponseException)).isTrue();
+  }
+
+  @Test
+  public void isHttpResponseExceptionRetriable_trueIf403WithRateLimitOverage() {
+    when(httpResponseException.getStatusCode()).thenReturn(403);
+    when(httpResponseException.getStatusMessage()).thenReturn("rateLimitExceeded");
+    assertThat(RetriableGmailExceptionPredicate.INSTANCE.test(httpResponseException)).isTrue();
+  }
+
+  @Test
+  public void isHttpResponseExceptionRetriable_trueIf403WithUserRateLimitOverage() {
+    when(httpResponseException.getStatusCode()).thenReturn(403);
+    when(httpResponseException.getStatusMessage()).thenReturn("userRateLimitExceeded");
+    assertThat(RetriableGmailExceptionPredicate.INSTANCE.test(httpResponseException)).isTrue();
+  }
+
+  @Test
+  public void isHttpResponseExceptionRetriable_falseIf403WithLongLastingOverage() {
+    when(httpResponseException.getStatusCode()).thenReturn(403);
+    when(httpResponseException.getStatusMessage()).thenReturn("dailyLimitExceeded");
+    assertThat(RetriableGmailExceptionPredicate.INSTANCE.test(httpResponseException)).isFalse();
+  }
+
+  @Test
+  public void isHttpResponseExceptionRetriable_falseIfBadRequest() {
+    when(httpResponseException.getStatusCode()).thenReturn(400);
+    assertThat(RetriableGmailExceptionPredicate.INSTANCE.test(httpResponseException)).isFalse();
+  }
+}