Testing: make the WalletFiller code a file-level singleton and fillTestWithCash an extension method. For Java users not much changes, the class is still called WalletFiller and the signature remains the same. Re-order some arguments to make it easier to use when accepting the defaults.

contracts/src/main/kotlin/com/r3corda/contracts/testing/WalletFiller.kt

@@ -1,47 +1,65 @@
+@file:JvmName("WalletFiller")
 package com.r3corda.contracts.testing
 
 import com.r3corda.contracts.cash.Cash
-import com.r3corda.core.contracts.Amount
-import com.r3corda.core.contracts.Issued
-import com.r3corda.core.contracts.TransactionType
+import com.r3corda.core.contracts.*
 import com.r3corda.core.crypto.Party
 import com.r3corda.core.node.ServiceHub
+import com.r3corda.core.node.services.Wallet
 import com.r3corda.core.serialization.OpaqueBytes
+import com.r3corda.core.testing.DUMMY_NOTARY
 import java.util.*
 
-object WalletFiller {
+
 /**
  * Creates a random set of between (by default) 3 and 10 cash states that add up to the given amount and adds them
-     * to the wallet.
+ * to the wallet. This is intended for unit tests.
  *
  * The cash is self issued with the current nodes identity, as fetched from the storage service. Thus it
  * would not be trusted by any sensible market participant and is effectively an IOU. If it had been issued by
  * the central bank, well ... that'd be a different story altogether.
+ *
+ * The service hub needs to provide at least a key management service and a storage service.
+ *
+ * @return a wallet object that represents the generated states (it will NOT be the full wallet from the service hub!)
  */
-    fun fillWithSomeTestCash(services: ServiceHub, notary: Party, howMuch: Amount<Currency>, atLeastThisManyStates: Int = 3,
-                             atMostThisManyStates: Int = 10, rng: Random = Random(),
-                             ref: OpaqueBytes = OpaqueBytes(ByteArray(1, { 0 }))) {
+fun ServiceHub.fillWithSomeTestCash(howMuch: Amount<Currency>,
+                                    notary: Party = DUMMY_NOTARY,
+                                    atLeastThisManyStates: Int = 3,
+                                    atMostThisManyStates: Int = 10,
+                                    rng: Random = Random(),
+                                    ref: OpaqueBytes = OpaqueBytes(ByteArray(1, { 0 }))): Wallet {
     val amounts = calculateRandomlySizedAmounts(howMuch, atLeastThisManyStates, atMostThisManyStates, rng)
 
-        val myIdentity = services.storageService.myLegalIdentity
-        val myKey = services.storageService.myLegalIdentityKey
+    val myIdentity = storageService.myLegalIdentity
+    val myKey = storageService.myLegalIdentityKey
 
     // We will allocate one state to one transaction, for simplicities sake.
     val cash = Cash()
-        val transactions = amounts.map { pennies ->
+    val transactions: List<SignedTransaction> = amounts.map { pennies ->
         // This line is what makes the cash self issued. We just use zero as our deposit reference: we don't need
         // this field as there's no other database or source of truth we need to sync with.
         val depositRef = myIdentity.ref(ref)
 
         val issuance = TransactionType.General.Builder()
-            val freshKey = services.keyManagementService.freshKey()
+        val freshKey = keyManagementService.freshKey()
         cash.generateIssue(issuance, Amount(pennies, Issued(depositRef, howMuch.token)), freshKey.public, notary)
         issuance.signWith(myKey)
 
         return@map issuance.toSignedTransaction(true)
     }
 
-        services.recordTransactions(transactions)
+    recordTransactions(transactions)
+
+    // Get all the StateRefs of all the generated transactions.
+    val states = transactions.flatMap { stx ->
+        stx.tx.outputs.indices.map { i -> stx.tx.outRef<Cash.State>(i) }
+    }
+
+    return object : Wallet() {
+        override val states: List<StateAndRef<ContractState>> = states
+        override val cashBalances: Map<Currency, Amount<Currency>> = mapOf(howMuch.token to howMuch)
+    }
 }
 
 private fun calculateRandomlySizedAmounts(howMuch: Amount<Currency>, min: Int, max: Int, rng: Random): LongArray {
@@ -61,4 +79,3 @@ object WalletFiller {
     }
     return amounts
 }
-}

core/src/main/kotlin/com/r3corda/core/node/services/Services.kt

@@ -67,7 +67,6 @@ abstract class Wallet {
     companion object {
         val NoUpdate = Update(emptySet(), emptySet())
     }
-
 }
 
 /**

node/src/main/kotlin/com/r3corda/node/internal/testing/TradeSimulation.kt

@@ -3,7 +3,7 @@ package com.r3corda.node.internal.testing
 import com.google.common.util.concurrent.Futures
 import com.google.common.util.concurrent.ListenableFuture
 import com.r3corda.contracts.CommercialPaper
-import com.r3corda.contracts.testing.WalletFiller
+import com.r3corda.contracts.testing.fillWithSomeTestCash
 import com.r3corda.core.contracts.DOLLARS
 import com.r3corda.core.contracts.SignedTransaction
 import com.r3corda.core.contracts.`issued by`
@@ -30,7 +30,7 @@ class TradeSimulation(runAsync: Boolean, latencyInjector: InMemoryMessagingNetwo
         val buyer = banks[buyerBankIndex]
         val seller = banks[sellerBankIndex]
 
-        WalletFiller.fillWithSomeTestCash(buyer.services, notary.info.identity, 1500.DOLLARS)
+        buyer.services.fillWithSomeTestCash(1500.DOLLARS, notary.info.identity)
 
         val issuance = run {
             val tx = CommercialPaper().generateIssue(1100.DOLLARS `issued by` seller.info.identity.ref(1, 2, 3), Instant.now() + 10.days, notary.info.identity)

node/src/test/kotlin/com/r3corda/node/messaging/TwoPartyTradeProtocolTests.kt

@@ -4,9 +4,9 @@ import com.google.common.util.concurrent.ListenableFuture
 import com.r3corda.contracts.CommercialPaper
 import com.r3corda.contracts.cash.Cash
 import com.r3corda.contracts.testing.CASH
-import com.r3corda.contracts.testing.WalletFiller
 import com.r3corda.contracts.testing.`issued by`
 import com.r3corda.contracts.testing.`owned by`
+import com.r3corda.contracts.testing.fillWithSomeTestCash
 import com.r3corda.core.contracts.*
 import com.r3corda.core.crypto.Party
 import com.r3corda.core.crypto.SecureHash
@@ -93,7 +93,7 @@ class TwoPartyTradeProtocolTests {
             val aliceNode = net.createPartyNode(notaryNode.info, ALICE.name, ALICE_KEY)
             val bobNode = net.createPartyNode(notaryNode.info, BOB.name, BOB_KEY)
 
-            WalletFiller.fillWithSomeTestCash(bobNode.services, DUMMY_NOTARY, 2000.DOLLARS)
+            bobNode.services.fillWithSomeTestCash(2000.DOLLARS)
             val issuer = bobNode.services.storageService.myLegalIdentity.ref(0)
             val alicesFakePaper = fillUpForSeller(false, aliceNode.storage.myLegalIdentity.owningKey,
                     1200.DOLLARS `issued by` issuer, notaryNode.info.identity, null).second
@@ -147,7 +147,7 @@ class TwoPartyTradeProtocolTests {
 
             net.runNetwork() // Clear network map registration messages
 
-            WalletFiller.fillWithSomeTestCash(bobNode.services, DUMMY_NOTARY, 2000.DOLLARS)
+            bobNode.services.fillWithSomeTestCash(2000.DOLLARS)
             val alicesFakePaper = fillUpForSeller(false, aliceNode.storage.myLegalIdentity.owningKey,
                     1200.DOLLARS `issued by` issuer, notaryNode.info.identity, null).second
             insertFakeTransactions(alicesFakePaper, aliceNode.services, aliceNode.storage.myLegalIdentityKey)

node/src/test/kotlin/com/r3corda/node/services/NodeWalletServiceTest.kt

@@ -1,7 +1,7 @@
 package com.r3corda.node.services
 
 import com.r3corda.contracts.cash.Cash
-import com.r3corda.contracts.testing.WalletFiller
+import com.r3corda.contracts.testing.fillWithSomeTestCash
 import com.r3corda.core.contracts.*
 import com.r3corda.core.node.ServiceHub
 import com.r3corda.core.node.services.testing.MockKeyManagementService
@@ -42,7 +42,7 @@ class NodeWalletServiceTest {
 
         kms.nextKeys += Array(3) { ALICE_KEY }
         // Fix the PRNG so that we get the same splits every time.
-        WalletFiller.fillWithSomeTestCash(services, DUMMY_NOTARY, 100.DOLLARS, 3, 3, Random(0L), ref)
+        services.fillWithSomeTestCash(100.DOLLARS, DUMMY_NOTARY, 3, 3, Random(0L), ref)
 
         val w = wallet.currentWallet
         assertEquals(3, w.states.size)

src/main/kotlin/com/r3corda/demos/TraderDemo.kt

@@ -3,7 +3,7 @@ package com.r3corda.demos
 import co.paralleluniverse.fibers.Suspendable
 import com.google.common.net.HostAndPort
 import com.r3corda.contracts.CommercialPaper
-import com.r3corda.contracts.testing.WalletFiller
+import com.r3corda.contracts.testing.fillWithSomeTestCash
 import com.r3corda.core.contracts.*
 import com.r3corda.core.crypto.Party
 import com.r3corda.core.crypto.SecureHash
@@ -226,7 +226,7 @@ class TraderDemoProtocolBuyer(private val attachmentsPath: Path,
         // Self issue some cash.
         //
         // TODO: At some point this demo should be extended to have a central bank node.
-        WalletFiller.fillWithSomeTestCash(serviceHub, notary, 3000.DOLLARS)
+        serviceHub.fillWithSomeTestCash(3000.DOLLARS, notary)
 
         while (true) {
             // Wait around until a node asks to start a trade with us. In a real system, this part would happen out of band