From 945aae816fc2ece38f9cadc1aa0a000385ee2fd9 Mon Sep 17 00:00:00 2001
From: jeremyhi <jiachun_fjc@163.com>
Date: Tue, 30 Jul 2024 10:21:26 +0800
Subject: [PATCH] feat: add hint api for Context

---
 ingester-rpc/pom.xml                          | 12 +++
 .../main/java/io/greptime/rpc/Context.java    | 24 +++++
 .../java/io/greptime/rpc/ContextTest.java     | 93 +++++++++++++++++++
 3 files changed, 129 insertions(+)
 create mode 100644 ingester-rpc/src/test/java/io/greptime/rpc/ContextTest.java

diff --git a/ingester-rpc/pom.xml b/ingester-rpc/pom.xml
index 874adba..9284055 100644
--- a/ingester-rpc/pom.xml
+++ b/ingester-rpc/pom.xml
@@ -31,5 +31,17 @@
             <groupId>${project.groupId}</groupId>
             <artifactId>ingester-common</artifactId>
         </dependency>
+
+        <!-- test -->
+        <dependency>
+            <groupId>junit</groupId>
+            <artifactId>junit</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-all</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 </project>
diff --git a/ingester-rpc/src/main/java/io/greptime/rpc/Context.java b/ingester-rpc/src/main/java/io/greptime/rpc/Context.java
index b057e92..c195c12 100644
--- a/ingester-rpc/src/main/java/io/greptime/rpc/Context.java
+++ b/ingester-rpc/src/main/java/io/greptime/rpc/Context.java
@@ -28,6 +28,8 @@
 @SuppressWarnings({"unchecked", "unused"})
 public class Context implements Copiable<Context> {
 
+    private static final String HINT_PREFIX = "x-greptime-hint:";
+
     private final Map<String, Object> ctx = new HashMap<>();
 
     /**
@@ -48,6 +50,17 @@ public static Context of(String key, Object value) {
         return new Context().with(key, value);
     }
 
+    /**
+     * Creates a new {@link Context} with the specified hint key-value pair.
+     *
+     * @param key the hint key
+     * @param value the value
+     * @return the new {@link Context}
+     */
+    public static Context hint(String key, Object value) {
+        return Context.of(HINT_PREFIX + key, value);
+    }
+
     /**
      * Adds the specified key-value pair to this {@link Context}.
      *
@@ -62,6 +75,17 @@ public Context with(String key, Object value) {
         return this;
     }
 
+    /**
+     * Adds the specified hint key-value pair to this {@link Context}.
+     *
+     * @param key the hint key
+     * @param value the value
+     * @return this {@link Context}
+     */
+    public Context withHint(String key, Object value) {
+        return with(HINT_PREFIX + key, value);
+    }
+
     /**
      * Gets the value of the specified key.
      *
diff --git a/ingester-rpc/src/test/java/io/greptime/rpc/ContextTest.java b/ingester-rpc/src/test/java/io/greptime/rpc/ContextTest.java
new file mode 100644
index 0000000..e2db367
--- /dev/null
+++ b/ingester-rpc/src/test/java/io/greptime/rpc/ContextTest.java
@@ -0,0 +1,93 @@
+/*
+ * Copyright 2023 Greptime Team
+ *
+ * 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.greptime.rpc;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+/**
+ * @author jiachun.fjc
+ */
+public class ContextTest {
+    @Test
+    public void newDefaultShouldReturnEmptyContextTest() {
+        Context context = Context.newDefault();
+        Assert.assertTrue(context.entrySet().isEmpty());
+    }
+
+    @Test
+    public void ofShouldReturnContextWithOneEntryTest() {
+        Context context = Context.of("key", "value");
+        Assert.assertEquals("value", context.get("key"));
+    }
+
+    @Test
+    public void hintShouldReturnContextWithHintEntryTest() {
+        Context context = Context.hint("key", "value");
+        Assert.assertEquals("value", context.get("x-greptime-hint:key"));
+    }
+
+    @Test
+    public void withShouldAddEntryToContextTest() {
+        Context context = Context.newDefault();
+        context.with("key", "value");
+        Assert.assertEquals("value", context.get("key"));
+    }
+
+    @Test
+    public void withHintShouldAddHintEntryToContextTest() {
+        Context context = Context.newDefault();
+        context.withHint("key", "value")
+                .with("key2", "value");
+        Assert.assertEquals("value", context.get("x-greptime-hint:key"));
+        Assert.assertEquals("value", context.get("key2"));
+    }
+
+    @Test
+    public void getShouldReturnNullForNonExistingKeyTest() {
+        Context context = Context.newDefault();
+        Assert.assertNull(context.get("key"));
+    }
+
+    @Test
+    public void removeShouldRemoveEntryFromContextTest() {
+        Context context = Context.of("key", "value");
+        String value = context.remove("key");
+        Assert.assertNull(context.get("key"));
+        Assert.assertEquals("value", value);
+    }
+
+    @Test
+    public void getOrDefaultShouldReturnDefaultValueForNonExistingKeyTest() {
+        Context context = Context.newDefault();
+        Assert.assertEquals("default", context.getOrDefault("key", "default"));
+    }
+
+    @Test
+    public void clearShouldRemoveAllEntriesFromContextTest() {
+        Context context = Context.of("key", "value");
+        context.clear();
+        Assert.assertTrue(context.entrySet().isEmpty());
+    }
+
+    @Test
+    public void copyShouldReturnIdenticalContextTest() {
+        Context context = Context.of("key", "value");
+        Context copy = context.copy();
+        Assert.assertEquals(context.get("key").toString(), copy.get("key").toString());
+    }
+}