Add registry test

mohsaka · Oct 14, 2024 · 8ddbc3d · 8ddbc3d
1 parent 61dacd0
commit 8ddbc3d
Showing 1 changed file with 35 additions and 0 deletions.
diff --git a/velox/functions/tests/FunctionRegistryTest.cpp b/velox/functions/tests/FunctionRegistryTest.cpp
@@ -26,6 +26,7 @@
 #include "velox/functions/Registerer.h"
 #include "velox/functions/prestosql/registration/RegistrationFunctions.h"
 #include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h"
+#include "velox/functions/prestosql/types/IPPrefixType.h"
 #include "velox/type/Type.h"
 
 namespace facebook::velox {
@@ -108,6 +109,17 @@ struct VariadicFunc {
   }
 };
 
+template <typename T>
+struct IPPrefixFunc {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE bool call(
+      out_type<IPPrefix>& /* result */,
+      const arg_type<IPPrefix>& /* arg1 */) {
+    return true;
+  }
+};
+
 class VectorFuncOne : public velox::exec::VectorFunction {
  public:
   void apply(
@@ -355,6 +367,14 @@ TEST_F(FunctionRegistryTest, getFunctionSignatures) {
           .build()
           ->toString());
 
+  ASSERT_EQ(
+      functionSignatures["ipprefix_func"].at(0)->toString(),
+      exec::FunctionSignatureBuilder()
+          .returnType("ipprefix")
+          .variableArity("ipprefix")
+          .build()
+          ->toString());
+
   ASSERT_EQ(
       functionSignatures["vector_func_one"].at(0)->toString(),
       exec::FunctionSignatureBuilder()
@@ -654,4 +674,19 @@ TEST_F(FunctionRegistryOverwriteTest, overwrite) {
   ASSERT_EQ(signatures.size(), 1);
 }
 
+TEST_F(FunctionRegistryTest, ipPrefixRegistration) {
+  registerFunction<IPPrefixFunc, IPPrefix, IPPrefix>({"ipprefix_func"});
+
+  auto& simpleFunctions = exec::simpleFunctions();
+  auto signatures = simpleFunctions.getFunctionSignatures("ipprefix_func");
+  ASSERT_EQ(signatures.size(), 1);
+
+  auto result = resolveFunctionWithMetadata("ipprefix_func", {IPPREFIX()});
+  EXPECT_TRUE(result.has_value());
+  EXPECT_EQ(*result->first, *IPPREFIX());
+  EXPECT_TRUE(result->second.defaultNullBehavior);
+  EXPECT_TRUE(result->second.deterministic);
+  EXPECT_FALSE(result->second.supportsFlattening);
+}
+
 } // namespace facebook::velox