Blackbox testing unifying abi-return-subroutine and subroutine input handling

pyteal/ast/router.py

@@ -133,7 +133,7 @@ class OnCompleteAction:
     def __post_init__(self):
         if bool(self.call_config) ^ bool(self.action):
             raise TealInputError(
-                f"action {self.action} and call_config {str(self.call_config)} contradicts"
+                f"action {self.action} and call_config {self.call_config!r} contradicts"
             )
 
     @staticmethod
@@ -219,7 +219,7 @@ def approval_construction(self) -> Optional[Expr]:
                     )
                 case _:
                     raise TealInternalError(
-                        f"Unexpected CallConfig: {str(oca.call_config)}"
+                        f"Unexpected CallConfig: {oca.call_config!r}"
                     )
             conditions_n_branches.append(
                 CondNode(
@@ -254,7 +254,7 @@ def clear_state_construction(self) -> Optional[Expr]:
                 )
             case _:
                 raise TealInternalError(
-                    f"Unexpected CallConfig: {str(self.clear_state.call_config)}"
+                    f"Unexpected CallConfig: {self.clear_state.call_config!r}"
                 )
 
 

tests/blackbox.py

@@ -70,9 +70,7 @@ def _fill(
         self, input_types: list[TealType | None]
     ) -> list[TealType | abi.TypeSpec | None]:
         match self.subroutine:
-            case SubroutineFnWrapper():
-                return cast(list[TealType | abi.TypeSpec | None], input_types)
-            case ABIReturnSubroutine():
+            case SubroutineFnWrapper() | ABIReturnSubroutine():
                 args = self.subroutine.subroutine.arguments()
                 abis = self.subroutine.subroutine.abi_args
                 return [(x if x else abis[args[i]]) for i, x in enumerate(input_types)]
@@ -123,6 +121,11 @@ def string_mult(x: abi.String, y: abi.Uint16):
     def string_mult(x: abi.String, y):
         ...
 
+    @Blackbox([None])
+    @Subroutine(TealType.uint64)
+    def cubed(n: abi.Uint64):
+        ...
+
     """
 
     def decorator_blackbox(func: SubroutineFnWrapper | ABIReturnSubroutine):
@@ -239,11 +242,11 @@ def program(self) -> Expr:
                 read arg and pass to subroutine as is
             * Expr arg of TealType.uint64:
                 convert arg to int using Btoi() when received
-            * pass-by-ref ScratchVar arguments (Subroutine case only):
+            * pass-by-ref ScratchVar arguments:
                 in addition to the above -
                     o store the arg (or converted arg) in a ScratchVar
                     o invoke the subroutine using this ScratchVar instead of the arg (or converted arg)
-            * ABI arguments (ABIReturnSubroutine case only):
+            * ABI arguments:
                 in addition to the above -
                     o store the decoded arg into the ScratchVar of an ABI Type instance
                     o invoke the subroutine using this ABI Type instead of the arg
@@ -282,31 +285,38 @@ def program(self) -> Expr:
 
         return self._pyteal_lambda()
 
-    def _handle_SubroutineFnWrapper(self):
+    def _arg_prep_n_call(self, i, p):
         subdef = self.subr.subroutine.subroutine
         arg_names = subdef.arguments()
-
-        def arg_prep_n_call(i, p):
-            name = arg_names[i]
-            by_ref = name in subdef.by_ref_args
-            arg_expr = (
-                Txn.application_args[i] if self.mode == Mode.Application else Arg(i)
-            )
-            if p == TealType.uint64:
-                arg_expr = Btoi(arg_expr)
-            prep = None
+        name = arg_names[i]
+        arg_expr = Txn.application_args[i] if self.mode == Mode.Application else Arg(i)
+        if p == TealType.uint64:
+            arg_expr = Btoi(arg_expr)
+
+        if name in subdef.by_ref_args:
+            arg_var = ScratchVar(p)
+            prep = arg_var.store(arg_expr)
+        elif name in subdef.abi_args:
+            arg_var = p.new_instance()
+            prep = arg_var.decode(arg_expr)
+        else:
             arg_var = arg_expr
-            if by_ref:
-                arg_var = ScratchVar(p)
-                prep = arg_var.store(arg_expr)
-            return prep, arg_var
+            prep = None
+        return prep, arg_var
+
+    def _prepare_n_calls(self):
+        preps_n_calls = [
+            *(self._arg_prep_n_call(i, p) for i, p in enumerate(self.input_types))
+        ]
+        preps, calls = zip(*preps_n_calls) if preps_n_calls else ([], [])
+        preps = [p for p in preps if p]
+        return preps, calls
+
+    def _handle_SubroutineFnWrapper(self):
+        subdef = self.subr.subroutine.subroutine
 
         def subr_caller():
-            preps_n_calls = [
-                *(arg_prep_n_call(i, p) for i, p in enumerate(self.input_types))
-            ]
-            preps, calls = zip(*preps_n_calls) if preps_n_calls else ([], [])
-            preps = [p for p in preps if p]
+            preps, calls = self._prepare_n_calls()
             invocation = self.subr(*calls)
             if preps:
                 return Seq(*(preps + [invocation]))
@@ -351,36 +361,12 @@ def approval():
         return approval
 
     def _handle_ABIReturnSubroutine(self):
-        subdef = self.subr.subroutine.subroutine
-        arg_names = subdef.arguments()
-
-        def arg_prep_n_call(i, p):
-            name = arg_names[i]
-            arg_expr = (
-                Txn.application_args[i] if self.mode == Mode.Application else Arg(i)
-            )
-            if p == TealType.uint64:
-                arg_expr = Btoi(arg_expr)
-            prep = None
-            arg_var = arg_expr
-            if name in subdef.by_ref_args:
-                arg_var = ScratchVar(p)
-                prep = arg_var.store(arg_expr)
-            elif name in subdef.abi_args:
-                arg_var = p.new_instance()
-                prep = arg_var.decode(arg_expr)
-            return prep, arg_var
-
         output = None
         if self.subr.subroutine.output_kwarg_info:
             output = self.subr.subroutine.output_kwarg_info.abi_type.new_instance()
 
         def approval():
-            preps_n_calls = [
-                *(arg_prep_n_call(i, p) for i, p in enumerate(self.input_types))
-            ]
-            preps, calls = zip(*preps_n_calls) if preps_n_calls else ([], [])
-            preps = [p for p in preps if p]
+            preps, calls = self._prepare_n_calls()
 
             # when @ABIReturnSubroutine is void:
             #   invocation is an Expr of TealType.none

tests/integration/graviton_test.py

@@ -829,6 +829,41 @@ def report(kind):
     report("lsig")
 
 
+def blackbox_pyteal_example5():
+    from graviton.blackbox import DryRunEncoder
+
+    from pyteal import abi, Subroutine, TealType, Int, Mode
+    from tests.blackbox import Blackbox
+
+    @Blackbox([None])
+    @Subroutine(TealType.uint64)
+    def cubed(n: abi.Uint64):
+        return n.get() ** Int(3)
+
+    app_pytealer = PyTealDryRunExecutor(cubed, Mode.Application)
+    lsig_pytealer = PyTealDryRunExecutor(cubed, Mode.Signature)
+
+    inputs = [[i] for i in range(1, 11)]
+
+    app_inspect = app_pytealer.dryrun_on_sequence(inputs)
+    lsig_inspect = lsig_pytealer.dryrun_on_sequence(inputs)
+
+    for index, inspect in enumerate(app_inspect):
+        input_var = inputs[index][0]
+        assert inspect.stack_top() == input_var**3, inspect.report(
+            args=inputs[index], msg="stack_top() gave unexpected results from app"
+        )
+        assert inspect.last_log() == DryRunEncoder.hex(input_var**3), inspect.report(
+            args=inputs[index], msg="last_log() gave unexpected results from app"
+        )
+
+    for index, inspect in enumerate(lsig_inspect):
+        input_var = inputs[index][0]
+        assert inspect.stack_top() == input_var**3, inspect.report(
+            args=inputs[index], msg="stack_top() gave unexpected results from app"
+        )
+
+
 def blackbox_pyteal_while_continue_test():
     from tests.blackbox import Blackbox
     from pyteal import (
@@ -880,6 +915,7 @@ def while_continue_accumulation(n):
         blackbox_pyteal_example2,
         blackbox_pyteal_example3,
         blackbox_pyteal_example4,
+        blackbox_pyteal_example5,
         blackbox_pyteal_while_continue_test,
     ],
 )