YT: Fix ParseEnum when custom domain names used

3 files changed, 39 insertions, 1 deletions

diff --git a/library/cpp/yt/misc/enum-inl.h b/library/cpp/yt/misc/enum-inl.h
index 951e93c130..c7e0abdab2 100644
--- a/library/cpp/yt/misc/enum-inl.h
+++ b/library/cpp/yt/misc/enum-inl.h
@@ -60,6 +60,23 @@ constexpr bool CheckValuesUnique(const TValues& values)
     return true;
 }
 
+template <typename TNames>
+constexpr bool CheckDomainNames(const TNames& names)
+{
+    for (size_t i = 0; i < std::size(names); ++i) {
+        if (std::size(names[i]) == 0) {
+            return false;
+        }
+        for (size_t j = 1; j < std::size(names[i]); ++j) {
+            // If name does not start with a capital letter, all the others must be in lowercase.
+            if (('A' <= names[i][j] && names[i][j] <= 'Z') && ('A' > names[i][0] || names[i][0] > 'Z')) {
+                return false;
+            }
+        }
+    }
+    return true;
+}
+
 } // namespace NDetail
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -81,6 +98,8 @@ constexpr bool CheckValuesUnique(const TValues& values)
         static constexpr std::array<TStringBuf, DomainSize> Names{{ \
             PP_FOR_EACH(ENUM__GET_DOMAIN_NAMES_ITEM, seq) \
         }}; \
+        static_assert(::NYT::NDetail::CheckDomainNames(Names), \
+            "Enumeration " #enumType " contains names in wrong format"); \
         static constexpr std::array<T, DomainSize> Values{{ \
             PP_FOR_EACH(ENUM__GET_DOMAIN_VALUES_ITEM, seq) \
         }}; \
diff --git a/library/cpp/yt/string/enum-inl.h b/library/cpp/yt/string/enum-inl.h
index 205ececae3..d9f8239cbd 100644
--- a/library/cpp/yt/string/enum-inl.h
+++ b/library/cpp/yt/string/enum-inl.h
@@ -33,7 +33,8 @@ std::optional<T> TryParseEnum(TStringBuf value)
 {
     auto tryFromString = [] (TStringBuf value) -> std::optional<T> {
         if (auto decodedValue = TryDecodeEnumValue(value)) {
-            return TEnumTraits<T>::FindValueByLiteral(*decodedValue);
+            auto enumValue = TEnumTraits<T>::FindValueByLiteral(*decodedValue);
+            return enumValue ? enumValue : TEnumTraits<T>::FindValueByLiteral(value);
         }
 
         auto reportError = [value] () {
diff --git a/library/cpp/yt/string/unittests/enum_ut.cpp b/library/cpp/yt/string/unittests/enum_ut.cpp
index fed6d6eb21..f30aec1cab 100644
--- a/library/cpp/yt/string/unittests/enum_ut.cpp
+++ b/library/cpp/yt/string/unittests/enum_ut.cpp
@@ -29,6 +29,11 @@ DEFINE_BIT_ENUM(ELangs,
     ((JavaScript) (0x10))
 );
 
+DEFINE_ENUM(ECustomDomainName,
+    ((A) (1) ("value_a"))
+    ((B) (2) ("value_b"))
+);
+
 TEST(TFormatTest, Enum)
 {
     EXPECT_EQ("Red", Format("%v", EColor::Red));
@@ -53,6 +58,19 @@ TEST(TFormatTest, Enum)
     EXPECT_EQ("cpp | go | python | java_script", Format("%lv", four));
 }
 
+TEST(TFormatEnumTest, FormatEnumWithCustomDomainName)
+{
+    EXPECT_EQ("value_a", FormatEnum(ECustomDomainName::A));
+    EXPECT_EQ("value_b", FormatEnum(ECustomDomainName::B));
+}
+
+TEST(TParseEnumTest, ParseEnumWithCustomDomainName)
+{
+    EXPECT_EQ(ECustomDomainName::A, TryParseEnum<ECustomDomainName>("value_a"));
+    EXPECT_EQ(ECustomDomainName::B, TryParseEnum<ECustomDomainName>("value_b"));
+    EXPECT_EQ(std::nullopt, TryParseEnum<ECustomDomainName>("b"));
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace