From 822db196c41c4305e33576eea350aec7e8b57661 Mon Sep 17 00:00:00 2001
From: Toshihiro Suzuki <brfrn169@gmail.com>
Date: Thu, 9 Jan 2025 14:27:04 +0900
Subject: [PATCH] Add utility setter methods for ImmutableMap to
 AbacOperationAttributes (#2448)

---
 .../db/api/AbacOperationAttributes.java       | 11 ++++++
 .../db/api/AbacOperationAttributesTest.java   | 35 +++++++++++++++++--
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/core/src/main/java/com/scalar/db/api/AbacOperationAttributes.java b/core/src/main/java/com/scalar/db/api/AbacOperationAttributes.java
index b558f8726..c4095f6a7 100644
--- a/core/src/main/java/com/scalar/db/api/AbacOperationAttributes.java
+++ b/core/src/main/java/com/scalar/db/api/AbacOperationAttributes.java
@@ -1,5 +1,6 @@
 package com.scalar.db.api;
 
+import com.google.common.collect.ImmutableMap;
 import java.util.Map;
 import java.util.Optional;
 
@@ -16,6 +17,11 @@ public static void setReadTag(Map<String, String> attributes, String policyName,
     attributes.put(READ_TAG_PREFIX + policyName, readTag);
   }
 
+  public static void setReadTag(
+      ImmutableMap.Builder<String, String> attributesBuilder, String policyName, String readTag) {
+    attributesBuilder.put(READ_TAG_PREFIX + policyName, readTag);
+  }
+
   public static void clearReadTag(Map<String, String> attributes, String policyName) {
     attributes.remove(READ_TAG_PREFIX + policyName);
   }
@@ -29,6 +35,11 @@ public static void setWriteTag(
     attributes.put(WRITE_TAG_PREFIX + policyName, writeTag);
   }
 
+  public static void setWriteTag(
+      ImmutableMap.Builder<String, String> attributesBuilder, String policyName, String writeTag) {
+    attributesBuilder.put(WRITE_TAG_PREFIX + policyName, writeTag);
+  }
+
   public static void clearWriteTag(Map<String, String> attributes, String policyName) {
     attributes.remove(WRITE_TAG_PREFIX + policyName);
   }
diff --git a/core/src/test/java/com/scalar/db/api/AbacOperationAttributesTest.java b/core/src/test/java/com/scalar/db/api/AbacOperationAttributesTest.java
index 56ed98af0..f58351234 100644
--- a/core/src/test/java/com/scalar/db/api/AbacOperationAttributesTest.java
+++ b/core/src/test/java/com/scalar/db/api/AbacOperationAttributesTest.java
@@ -2,6 +2,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
+import com.google.common.collect.ImmutableMap;
 import com.scalar.db.io.Key;
 import java.util.HashMap;
 import java.util.Map;
@@ -11,7 +12,7 @@
 public class AbacOperationAttributesTest {
 
   @Test
-  public void setReadTag_ShouldSetReadTag() {
+  public void setReadTag_MapGiven_ShouldSetReadTag() {
     // Arrange
     Map<String, String> attributes = new HashMap<>();
     String policyName = "policyName";
@@ -25,6 +26,21 @@ public void setReadTag_ShouldSetReadTag() {
         .containsEntry(AbacOperationAttributes.READ_TAG_PREFIX + policyName, readTag);
   }
 
+  @Test
+  public void setReadTag_ImmutableMapBuilderGiven_ShouldSetReadTag() {
+    // Arrange
+    ImmutableMap.Builder<String, String> attributesBuilder = ImmutableMap.builder();
+    String policyName = "policyName";
+    String readTag = "readTag";
+
+    // Act
+    AbacOperationAttributes.setReadTag(attributesBuilder, policyName, readTag);
+
+    // Assert
+    assertThat(attributesBuilder.build())
+        .containsEntry(AbacOperationAttributes.READ_TAG_PREFIX + policyName, readTag);
+  }
+
   @Test
   public void clearReadTag_ShouldClearReadTag() {
     // Arrange
@@ -60,7 +76,7 @@ public void clearReadTags_ShouldClearReadTags() {
   }
 
   @Test
-  public void setWriteTag_ShouldSetWriteTag() {
+  public void setWriteTag_MapGiven_ShouldSetWriteTag() {
     // Arrange
     Map<String, String> attributes = new HashMap<>();
     String policyName = "policyName";
@@ -74,6 +90,21 @@ public void setWriteTag_ShouldSetWriteTag() {
         .containsEntry(AbacOperationAttributes.WRITE_TAG_PREFIX + policyName, writeTag);
   }
 
+  @Test
+  public void setWriteTag_ImmutableMapBuilderGiven_ShouldSetWriteTag() {
+    // Arrange
+    ImmutableMap.Builder<String, String> attributesBuilder = ImmutableMap.builder();
+    String policyName = "policyName";
+    String writeTag = "writeTag";
+
+    // Act
+    AbacOperationAttributes.setWriteTag(attributesBuilder, policyName, writeTag);
+
+    // Assert
+    assertThat(attributesBuilder.build())
+        .containsEntry(AbacOperationAttributes.WRITE_TAG_PREFIX + policyName, writeTag);
+  }
+
   @Test
   public void clearWriteTag_ShouldClearWriteTag() {
     // Arrange