From 4871d55b5605f6a302d75e7a20c4bc0f8d8ab5f3 Mon Sep 17 00:00:00 2001
From: yujingwei <yujingwei@sensorsdata.cn>
Date: Wed, 18 Oct 2023 04:05:55 +0000
Subject: [PATCH] Add unit test and fix IWYU

---
 src/base/pegasus_utils.cpp                    |  3 -
 .../test/redact_sensitive_string_test.cpp     | 73 +++++++++++++++++++
 2 files changed, 73 insertions(+), 3 deletions(-)
 create mode 100644 src/base/test/redact_sensitive_string_test.cpp

diff --git a/src/base/pegasus_utils.cpp b/src/base/pegasus_utils.cpp
index 72b831a633..e0859fe3d1 100644
--- a/src/base/pegasus_utils.cpp
+++ b/src/base/pegasus_utils.cpp
@@ -22,14 +22,11 @@
 #include <arpa/inet.h>
 #include <netdb.h>
 #include <netinet/in.h>
-#include <rocksdb/slice.h>
 #include <stdio.h>
 #include <sys/socket.h>
 #include <cctype>
-#include <utility>
 
 #include "runtime/rpc/rpc_address.h"
-#include "utils/blob.h"
 #include "utils/fmt_logging.h"
 
 namespace pegasus {
diff --git a/src/base/test/redact_sensitive_string_test.cpp b/src/base/test/redact_sensitive_string_test.cpp
new file mode 100644
index 0000000000..1728f27a44
--- /dev/null
+++ b/src/base/test/redact_sensitive_string_test.cpp
@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ */
+
+// IWYU pragma: no_include <gtest/gtest-message.h>
+// IWYU pragma: no_include <gtest/gtest-test-part.h>
+#include <gtest/gtest.h>
+#include <list>
+#include <string>
+#include <iostream>
+
+#include "base/pegasus_utils.h"
+
+const std::string test_string = "pegasus";
+
+// FLAGS_encrypt_data_at_rest up
+TEST(pegasus_utils, redact_sensitive_string)
+{
+    FLAGS_encrypt_data_at_rest = true;
+    auto result_string = pegasus::utils::redact_sensitive_string(test_string);
+    ASSERT_EQ(pegasus::utils::kRedactedString, result_string);
+}
+
+// FLAGS_encrypt_data_at_rest down
+TEST(pegasus_utils, redact_sensitive_string_with_encrypt)
+{
+    FLAGS_encrypt_data_at_rest = false;
+    auto result_string = pegasus::utils::redact_sensitive_string(test_string);
+    ASSERT_EQ("pegasus", result_string);
+}
+
+TEST(pegasus_utils, c_escape_sensitive_string_with_no_encrypt_and_escape)
+{
+    FLAGS_encrypt_data_at_rest = false;
+    auto result_string = pegasus::utils::c_escape_sensitive_string(test_string, false);
+    ASSERT_EQ("pegasus", result_string);
+}
+
+TEST(pegasus_utils, c_escape_sensitive_string_with_no_encrypt)
+{
+    FLAGS_encrypt_data_at_rest = false;
+    auto result_string = pegasus::utils::c_escape_sensitive_string(test_string, true);
+    ASSERT_EQ("\\x70\\x65\\x67\\x61\\x73\\x75\\x73", result_string);
+}
+
+TEST(pegasus_utils, c_escape_sensitive_string_with_encrypt)
+{
+    FLAGS_encrypt_data_at_rest = true;
+    auto result_string = pegasus::utils::c_escape_sensitive_string(test_string, false);
+    ASSERT_EQ(pegasus::utils::kRedactedString, result_string);
+}
+
+TEST(pegasus_utils, c_escape_sensitive_string_with_encrypt_and_escape)
+{
+    FLAGS_encrypt_data_at_rest = true;
+    auto result_string = pegasus::utils::c_escape_sensitive_string(test_string, true);
+    ASSERT_EQ(pegasus::utils::kRedactedString, result_string);
+}