Clean ut

src/Neo.SmartContract.Testing/TestingStandards/Nep17Tests.cs

@@ -204,28 +204,37 @@ public virtual void TestTransfer()
 
         UInt160? calledFrom = null;
         BigInteger? calledAmount = null;
-        byte[]? calledData = null;
+        UInt160? calledData = null;
 
         var mock = Engine.Deploy<onNEP17PaymentContract>(NefFile, manifest.ToJson().ToString(), null, m =>
          {
              m
              .Setup(s => s.onNEP17Payment(It.IsAny<UInt160>(), It.IsAny<BigInteger>(), It.IsAny<object>()))
              .Callback(new InvocationAction((i) =>
              {
+                 // Set variables
+
+                 var me = new UInt160((i.Arguments[2] as ByteString)!.GetSpan().ToArray());
                  calledFrom = i.Arguments[0] as UInt160;
                  calledAmount = (BigInteger)i.Arguments[1];
-                 calledData = (i.Arguments[2] as ByteString)!.GetSpan().ToArray();
+
+                 // Ensure the balance
+
+                 Assert.AreEqual(3, Contract.BalanceOf(me));
 
                  // Ensure the event was called
 
-                 var me = new UInt160(calledData);
                  AssertTransferEvent(Alice.Account, me, calledAmount);
 
                  // Return the money back
 
                  Engine.SetTransactionSigners(me);
                  Assert.IsTrue(Contract.Transfer(me, calledFrom, calledAmount));
                  AssertTransferEvent(me, Alice.Account, calledAmount);
+
+                 // Set success flag
+
+                 calledData = me;
              }));
          });
 
@@ -235,7 +244,7 @@ public virtual void TestTransfer()
         Assert.IsTrue(Contract.Transfer(Alice.Account, mock.Hash, 3, mock.Hash.ToArray()));
 
         Assert.AreEqual(Alice.Account, calledFrom);
-        Assert.AreEqual(mock.Hash, new UInt160(calledData));
+        Assert.AreEqual(mock.Hash, calledData);
         Assert.AreEqual(3, calledAmount);
         Assert.AreEqual(0, Contract.BalanceOf(mock.Hash));
     }