Updated tests to use getBalance

rafaelspinto · Jan 7, 2019 · 93b5908 · 93b5908
1 parent 57edf92
commit 93b5908
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/src/main/scala/workshop/payment/CustomWallet.scala b/src/main/scala/workshop/payment/CustomWallet.scala
@@ -1,12 +1,12 @@
 package workshop.payment
 
-case class CustomWallet(balance: Int, maxAmount: Int = 1000) extends Wallet{
+case class CustomWallet(var balance: Int, maxAmount: Int = 1000) extends Wallet{
   def remove(amount: Int): Int = {
     if(amount > balance) {
       throw new NotEnoughFundsException
     }
-
-    balance - amount
+    balance -= amount
+    balance
   }
 
   def add(amount: Int): Int = {
@@ -16,7 +16,8 @@ case class CustomWallet(balance: Int, maxAmount: Int = 1000) extends Wallet{
     if(amount > maxAmount) {
       throw new MaxAmountExceededException
     }
-    amount + balance
+    balance += amount
+    balance
   }
 
   override def getBalance: Int = balance

diff --git a/src/test/scala/workshop/payment/CustomWalletTest.scala b/src/test/scala/workshop/payment/CustomWalletTest.scala
@@ -22,7 +22,7 @@ class CustomWalletTest extends FunSuite {
     val amount = 10
     val balance = 11
     val wallet: CustomWallet = new CustomWallet(balance)
-    assert(wallet.add(amount) == balance + amount)
+    assert(wallet.add(amount) == wallet.getBalance)
   }
 
   test("add: if amount is greater than allowed throws MaxAmountExceededException") {
@@ -45,6 +45,6 @@ class CustomWalletTest extends FunSuite {
     val amount = 9
     val balance = 11
     val wallet: CustomWallet = new CustomWallet(balance)
-    assert(wallet.remove(amount) == balance - amount)
+    assert(wallet.remove(amount) == wallet.getBalance)
   }
 }