fix unit tests

Author: chelsea-lin

bigframes/core/compile/sqlglot/aggregations/op_registration.py

@@ -41,21 +41,16 @@ def arg_checker(*args, **kwargs):
                     )
                 return item(*args, **kwargs)
 
-            key = op if isinstance(op, type) else type(op)
+            key = str(op)
             if key in self._registered_ops:
                 raise ValueError(f"{key} is already registered")
-            self._registered_ops[str(key)] = item
+            self._registered_ops[key] = item
             return arg_checker
 
         return decorator
 
     def __getitem__(self, op: str | agg_ops.WindowOp) -> CompilationFunc:
-        if isinstance(op, agg_ops.WindowOp):
-            if not hasattr(op, "name"):
-                raise ValueError(f"The operator must have a 'name' attribute. Got {op}")
-            else:
-                key = typing.cast(str, op.name)
-                if key in self._registered_ops:
-                    return self._registered_ops[key]
-                return self._registered_ops[str(type(op))]
-        return self._registered_ops[op]
+        key = op if isinstance(op, type) else type(op)
+        if str(key) not in self._registered_ops:
+            raise ValueError(f"{key} is already not registered")
+        return self._registered_ops[str(key)]

tests/unit/core/compile/sqlglot/aggregations/test_op_registration.py

@@ -29,7 +29,6 @@ def test_func(op: agg_ops.SizeOp, input: sge.Expression) -> sge.Expression:
         return input
 
     assert reg[agg_ops.SizeOp()](op, input) == test_func(op, input)
-    assert reg[agg_ops.SizeOp.name](op, input) == test_func(op, input)
 
 
 def test_register_function_first_argument_is_not_agg_op_raise_error():