From c1eba6d161079a118a1a7e282ffa2debdf28d743 Mon Sep 17 00:00:00 2001
From: Chong Shao <chongshao@google.com>
Date: Mon, 18 Nov 2019 11:01:07 -0800
Subject: [PATCH] Support direct_boot_ok parameter

---
 .../firebase/messaging/AndroidConfig.java     | 14 +++++++++++
 .../firebase/messaging/MessageTest.java       | 23 +++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/src/main/java/com/google/firebase/messaging/AndroidConfig.java b/src/main/java/com/google/firebase/messaging/AndroidConfig.java
index dc0a45526..d16190881 100644
--- a/src/main/java/com/google/firebase/messaging/AndroidConfig.java
+++ b/src/main/java/com/google/firebase/messaging/AndroidConfig.java
@@ -51,6 +51,9 @@ public class AndroidConfig {
 
   @Key("fcm_options")
   private final AndroidFcmOptions fcmOptions;
+  
+  @Key("direct_boot_ok")
+  private final Boolean directBootOk;
 
   private AndroidConfig(Builder builder) {
     this.collapseKey = builder.collapseKey;
@@ -75,6 +78,7 @@ private AndroidConfig(Builder builder) {
     this.data = builder.data.isEmpty() ? null : ImmutableMap.copyOf(builder.data);
     this.notification = builder.notification;
     this.fcmOptions = builder.fcmOptions;
+    this.directBootOk = builder.directBootOk;
   }
 
   /**
@@ -103,6 +107,7 @@ public static class Builder {
     private final Map<String, String> data = new HashMap<>();
     private AndroidNotification notification;
     private AndroidFcmOptions fcmOptions;
+    private Boolean directBootOk;
 
     private Builder() {}
 
@@ -201,6 +206,15 @@ public Builder setFcmOptions(AndroidFcmOptions androidFcmOptions) {
       return this;
     }
 
+    /**
+     * Sets the direct_boot_ok flag, If set to true, messages will be allowed to be delivered to 
+     * the app while the device is in direct boot mode.
+     */
+    public Builder setDirectBootOk(boolean directBootOk) {
+      this.directBootOk = directBootOk;
+      return this;
+    }
+
     /**
      * Creates a new {@link AndroidConfig} instance from the parameters set on this builder.
      *
diff --git a/src/test/java/com/google/firebase/messaging/MessageTest.java b/src/test/java/com/google/firebase/messaging/MessageTest.java
index 5c8c3cbbf..170760a1a 100644
--- a/src/test/java/com/google/firebase/messaging/MessageTest.java
+++ b/src/test/java/com/google/firebase/messaging/MessageTest.java
@@ -200,6 +200,29 @@ public void testAndroidMessageWithNotification() throws IOException {
     assertJsonEquals(ImmutableMap.of("topic", "test-topic", "android", data), message);
   }
 
+  @Test
+  public void testAndroidMessageWithDirectBootOk() throws IOException {
+    Message message = Message.builder()
+        .setAndroidConfig(AndroidConfig.builder()
+            .setDirectBootOk(true)
+            .setNotification(AndroidNotification.builder()
+                .setTitle("android-title")
+                .setBody("android-body")
+                .build())
+            .build())
+        .setTopic("test-topic")
+        .build();
+    Map<String, Object> notification = ImmutableMap.<String, Object>builder()
+        .put("title", "android-title")
+        .put("body", "android-body")
+        .build();
+    Map<String, Object> data = ImmutableMap.of(
+        "direct_boot_ok", true,
+        "notification", notification
+    );
+    assertJsonEquals(ImmutableMap.of("topic", "test-topic", "android", data), message);
+  }
+
   @Test(expected = IllegalArgumentException.class)
   public void testAndroidNotificationWithNegativeCount() throws IllegalArgumentException {
     AndroidNotification.builder().setNotificationCount(-1).build();