From 40b4782f9d1eaa5f4b29bb885134c607db8181a5 Mon Sep 17 00:00:00 2001
From: Mark Spanbroek <mark@spanbroek.net>
Date: Mon, 22 Mar 2021 09:23:33 +0100
Subject: [PATCH] Return updated signed state after payment

---
 nitro/wallet/wallet.nim    | 12 ++++++------
 tests/nitro/testWallet.nim |  7 +++++++
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/nitro/wallet/wallet.nim b/nitro/wallet/wallet.nim
index a41dacc..8f5ad3f 100644
--- a/nitro/wallet/wallet.nim
+++ b/nitro/wallet/wallet.nim
@@ -108,24 +108,24 @@ func pay*(wallet: var Wallet,
           channel: ChannelId,
           asset: EthAddress,
           receiver: Destination,
-          amount: UInt256): ?!void =
+          amount: UInt256): ?!SignedState =
   if var state =? wallet.state(channel):
     if var balances =? state.outcome.balances(asset):
       ?balances.move(wallet.destination, receiver, amount)
       try:
         state.outcome.update(asset, balances)
         wallet.updateChannel(SignedState(state: state))
-        ok()
+        ok(wallet.channels[channel])
       except KeyError as error:
-        void.failure error
+        SignedState.failure error
     else:
-      void.failure "asset not found"
+      SignedState.failure "asset not found"
   else:
-    void.failure "channel not found"
+    SignedState.failure "channel not found"
 
 func pay*(wallet: var Wallet,
           channel: ChannelId,
           asset: EthAddress,
           receiver: EthAddress,
-          amount: UInt256): ?!void =
+          amount: UInt256): ?!SignedState =
   wallet.pay(channel, asset, receiver.toDestination, amount)
diff --git a/tests/nitro/testWallet.nim b/tests/nitro/testWallet.nim
index fad378a..5b38281 100644
--- a/tests/nitro/testWallet.nim
+++ b/tests/nitro/testWallet.nim
@@ -106,6 +106,13 @@ suite "wallet: making payments":
     let expectedSignature = key.sign(wallet.state(channel).get)
     check wallet.signature(channel, wallet.address) == expectedSignature.some
 
+  test "pay returns the updated signed state":
+    wallet = Wallet.init(key)
+    channel = wallet.openLedgerChannel(hub, chainId, nonce, asset, 42.u256)
+    let updated = wallet.pay(channel, asset, hub, 1.u256).option
+    check updated?.state == wallet.state(channel)
+    check updated?.signatures == wallet.signatures(channel)
+
   test "payment fails when channel not found":
     wallet = Wallet.init(key)
     check wallet.pay(channel, asset, hub, 1.u256).isErr