improve PullPaymentCapable

contracts/PullPayment.sol

@@ -13,11 +13,16 @@ contract PullPayment {
   }
 
   // withdraw accumulated balance, called by payee
-  function withdrawPayments() external {
+  function withdrawPayments() {
-    uint payment = payments[msg.sender];
+    address payee = msg.sender;
-    payments[msg.sender] = 0;
+    uint payment = payments[payee];
-    if (!msg.sender.send(payment)) {
+    
-      payments[msg.sender] = payment;
+    if (payment == 0) throw;
+    if (this.balance < payment) throw;
+
+    payments[payee] = 0;
+    if (!payee.send(payment)) {
+      throw;
     }
   }
 }

contracts/examples/PullPaymentExample.sol

@@ -1,11 +0,0 @@
-pragma solidity ^0.4.0;
-
-import '../PullPayment.sol';
-
-// Example class using PullPayment
-contract PullPaymentExample is PullPayment {
-  // test helper function to call asyncSend
-  function callSend(address dest, uint amount) external {
-    asyncSend(dest, amount);
-  }
-}

contracts/test-helpers/PullPaymentMock.sol

@@ -4,7 +4,8 @@ import '../PullPayment.sol';
 // mock class using PullPayment
 contract PullPaymentMock is PullPayment {
   // test helper function to call asyncSend
-  function callSend(address dest, uint amount) external {
+  function callSend(address dest, uint amount) {
     asyncSend(dest, amount);
   }
+
 }

test/PullPayment.js

@@ -1,8 +1,7 @@
-contract('PullPaymentExample', function(accounts) {
+contract('PullPayment', function(accounts) {
 
   it("can't call asyncSend externally", function(done) {
-    var ppc;
+    return PullPaymentMock.new()
-    return PullPaymentExample.new()
       .then(function(ppc) {
         assert.isUndefined(ppc.asyncSend);
       })
@@ -12,7 +11,7 @@ contract('PullPaymentExample', function(accounts) {
   it("can record an async payment correctly", function(done) {
     var ppce;
     var AMOUNT = 100;
-    return PullPaymentExample.new()
+    return PullPaymentMock.new()
       .then(function(_ppce) {
         ppce = _ppce;
         ppce.callSend(accounts[0], AMOUNT)
@@ -26,4 +25,78 @@ contract('PullPaymentExample', function(accounts) {
       .then(done);
   });
 
+  it("can add multiple balances on one account", function(done) {
+    var ppce;
+    return PullPaymentMock.new()
+      .then(function(_ppce) {
+        ppce = _ppce;
+        return ppce.callSend(accounts[0], 200)
+      })
+      .then(function() {
+        return ppce.callSend(accounts[0], 300)
+      })
+      .then(function() {
+        return ppce.payments(accounts[0]);
+      })
+      .then(function(paymentsToAccount0) {
+        assert.equal(paymentsToAccount0, 500);
+      })
+      .then(done);
+  });
+
+  it("can add balances on multiple accounts", function(done) {
+    var ppce;
+    return PullPaymentMock.new()
+      .then(function(_ppce) {
+        ppce = _ppce;
+        return ppce.callSend(accounts[0], 200)
+      })
+      .then(function() {
+        return ppce.callSend(accounts[1], 300)
+      })
+      .then(function() {
+        return ppce.payments(accounts[0]);
+      })
+      .then(function(paymentsToAccount0) {
+        assert.equal(paymentsToAccount0, 200);
+      })
+      .then(function() {
+        return ppce.payments(accounts[1]);
+      })
+      .then(function(paymentsToAccount0) {
+        assert.equal(paymentsToAccount0, 300);
+      })
+      .then(done);
+  });
+
+  it("can withdraw payment", function(done) {
+    var ppce;
+    var AMOUNT = 17*1e18;
+    var payee = accounts[1];
+    var initialBalance = web3.eth.getBalance(payee);
+    return PullPaymentMock.new({value: AMOUNT})
+      .then(function(_ppce) {
+        ppce = _ppce;
+        return ppce.callSend(payee, AMOUNT);
+      })
+      .then(function() {
+        return ppce.payments(payee);
+      })
+      .then(function(paymentsToAccount0) {
+        assert.equal(paymentsToAccount0, AMOUNT);
+      })
+      .then(function() {
+        return ppce.withdrawPayments({from: payee});
+      })
+      .then(function() {
+        return ppce.payments(payee);
+      })
+      .then(function(paymentsToAccount0) {
+        assert.equal(paymentsToAccount0, 0);
+        var balance = web3.eth.getBalance(payee);
+        assert(Math.abs(balance-initialBalance-AMOUNT) < 1e16);
+      })
+      .then(done);
+  });
+
 });