From 6189b28489e5b7286117717ecc4a36f37d202955 Mon Sep 17 00:00:00 2001
From: Larry Luo <larry.luo@ca.ibm.com>
Date: Mon, 18 Apr 2022 11:34:25 -0400
Subject: [PATCH] Added unit tests for parsing sha256_password and sha256_hash
 with salt

---
 src/Parsers/tests/gtest_Parser.cpp | 56 ++++++++++++++++++++++++++++--
 1 file changed, 53 insertions(+), 3 deletions(-)

diff --git a/src/Parsers/tests/gtest_Parser.cpp b/src/Parsers/tests/gtest_Parser.cpp
index fac79de4c5b..d085ac23c67 100644
--- a/src/Parsers/tests/gtest_Parser.cpp
+++ b/src/Parsers/tests/gtest_Parser.cpp
@@ -5,13 +5,14 @@
 #include <Parsers/ASTIdentifier.h>
 #include <Parsers/ParserAlterQuery.h>
 #include <Parsers/ParserCreateQuery.h>
+#include <Parsers/Access/ParserCreateUserQuery.h>
 #include <Parsers/ParserOptimizeQuery.h>
 #include <Parsers/ParserQueryWithOutput.h>
 #include <Parsers/formatAST.h>
 #include <Parsers/parseQuery.h>
-
+#include "Access/AccessEntityIO.cpp"
 #include <string_view>
-
+#include <regex>
 #include <gtest/gtest.h>
 
 namespace
@@ -50,7 +51,24 @@ TEST_P(ParserTest, parseQuery)
     {
         ASTPtr ast;
         ASSERT_NO_THROW(ast = parseQuery(*parser, input_text.begin(), input_text.end(), 0, 0));
-        EXPECT_EQ(expected_ast, serializeAST(*ast->clone(), false));
+
+        if (std::string("CREATE USER or ALTER USER query").compare(parser->getName()) 
+                && std::string("ATTACH access entity query").compare(parser->getName()))
+        {
+            EXPECT_EQ(expected_ast, serializeAST(*ast->clone(), false));
+        }
+        else
+        {
+            if (input_text.starts_with("ATTACH"))
+            {
+                auto salt = (dynamic_cast<const ASTCreateUserQuery *>(ast.get())->auth_data)->getSalt();
+                EXPECT_TRUE(std::regex_match(salt, std::regex(expected_ast)));
+            }
+            else
+            {
+                EXPECT_TRUE(std::regex_match(serializeAST(*ast->clone(), false), std::regex(expected_ast)));
+            }
+        }
     }
     else
     {
@@ -226,3 +244,35 @@ INSTANTIATE_TEST_SUITE_P(ParserCreateDatabaseQuery, ParserTest,
             "CREATE DATABASE db\nENGINE = Foo\nSETTINGS a = 1, b = 2\nTABLE OVERRIDE `a`\n(\n    ORDER BY (`id`, `version`)\n)\nCOMMENT 'db comment'"
         }
 })));
+
+INSTANTIATE_TEST_SUITE_P(ParserCreateUserQuery, ParserTest,
+    ::testing::Combine(
+        ::testing::Values(std::make_shared<ParserCreateUserQuery>()),
+        ::testing::ValuesIn(std::initializer_list<ParserTestCase>{
+        {
+            "CREATE USER user1 IDENTIFIED WITH sha256_password BY 'qwe123'",
+            "CREATE USER user1 IDENTIFIED WITH sha256_hash BY '[A-Za-z0-9]{64}' '[A-Za-z0-9]{64}'"
+        },
+        {
+            "ALTER USER user1 IDENTIFIED WITH sha256_password BY 'qwe123'",
+            "ALTER USER user1 IDENTIFIED WITH sha256_hash BY '[A-Za-z0-9]{64}' '[A-Za-z0-9]{64}'"
+        },
+        {
+            "ATTACH USER user1 IDENTIFIED WITH sha256_hash BY '2CC4880302693485717D34E06046594CFDFE425E3F04AA5A094C4AABAB3CB0BF' 'EFFD7F6B03B3EA68B8F86C1E91614DD50E42EB31EF7160524916444D58B5E264';",
+            "ALTER USER user1 IDENTIFIED WITH sha256_hash BY '[A-Za-z0-9]{64}' '[A-Za-z0-9]{64}'"
+        }
+})));
+
+INSTANTIATE_TEST_SUITE_P(ParserAttachUserQuery, ParserTest,
+    ::testing::Combine(
+        ::testing::Values(std::make_shared<ParserAttachAccessEntity>()),
+        ::testing::ValuesIn(std::initializer_list<ParserTestCase>{
+        {
+            "ATTACH USER user1 IDENTIFIED WITH sha256_hash BY '2CC4880302693485717D34E06046594CFDFE425E3F04AA5A094C4AABAB3CB0BF' 'EFFD7F6B03B3EA68B8F86C1E91614DD50E42EB31EF7160524916444D58B5E264';",
+            "^[A-Za-z0-9]{64}$"
+        },
+        {
+            "ATTACH USER user1 IDENTIFIED WITH sha256_hash BY '2CC4880302693485717D34E06046594CFDFE425E3F04AA5A094C4AABAB3CB0BF'",  //for users created in older releases that sha256_password has no salt
+            "^$"
+        }
+})));