1 files changed, 93 insertions, 4 deletions

diff --git a/pbt/ts/src/bank.test.ts b/pbt/ts/src/bank.test.ts
index 9ec25c0..c448242 100644
--- a/pbt/ts/src/bank.test.ts
+++ b/pbt/ts/src/bank.test.ts
@@ -1,16 +1,18 @@
 import {toNullable} from 'fp-ts/lib/Option';
-import {Account, SimpleBank} from './bank';
+import {Account, Operation, SimpleBank} from './bank';
 import {
   Gen,
   Property,
   choose,
   elements,
   forall,
+  frequency,
   keyValueOf,
-  mapOf,
-  mapOf1,
+  listOf,
+  mapOf1 as nonEmptyMap,
   property,
 } from './property';
+import Prando from 'prando';
 
 const sampleAccounts: Account[] = [
   'FR12345678901234',
@@ -22,7 +24,81 @@ const genAccount: Gen<Account> = elements(sampleAccounts);
 
 const genBalance: Gen<number> = choose(0, 1000000);
 
-const genAccounts: Gen<Map<Account, number>> = mapOf1(genAccount, genBalance);
+const genAccounts: Gen<Map<Account, number>> = nonEmptyMap(
+  genAccount,
+  genBalance,
+);
+
+const genGetBalance: (accounts: Map<Account, number>) => Gen<Operation> =
+  (accounts: Map<Account, number>) => (rng: Prando) => (size: number) => {
+    const account = elements(Array.from(accounts.keys()))(rng)(size);
+    return {op: 'GetBalance', account};
+  };
+
+const genDeposit: (accounts: Map<Account, number>) => Gen<Operation> =
+  (accounts: Map<Account, number>) => (rng: Prando) => (size: number) => {
+    const account = elements(Array.from(accounts.keys()))(rng)(size);
+    const gAmt = choose(1, 100000)(rng)(size);
+    return {op: 'Deposit', account, amount: gAmt};
+  };
+
+const genWithdraw: (accounts: Map<Account, number>) => Gen<Operation> =
+  (accounts: Map<Account, number>) => (rng: Prando) => (size: number) => {
+    const account = elements(Array.from(accounts.keys()))(rng)(size);
+    const gAmt = choose(1, 100000)(rng)(size);
+    return {op: 'Withdraw', account, amount: gAmt};
+  };
+
+const genOneOperation = (accounts: Map<Account, number>) =>
+  frequency([
+    [1, genGetBalance(accounts)],
+    [3, genDeposit(accounts)],
+    [3, genWithdraw(accounts)],
+  ]);
+
+const genOperations = (accounts: Map<Account, number>) =>
+  listOf(genOneOperation(accounts));
+
+const correctBalance =
+  (api: SimpleBank, initialAccounts: Map<Account, number>) =>
+  (ops: Operation[]): boolean => {
+    const expectedBalances = new Map(initialAccounts);
+
+    for (const op of ops) {
+      switch (op.op) {
+        case 'GetBalance': {
+          const retrieved = toNullable(api.balance(op.account));
+          const expected = expectedBalances.get(op.account);
+          if (retrieved !== expected) {
+            return false;
+          }
+          break;
+        }
+        case 'Deposit': {
+          api.deposit(op.account, op.amount);
+          const current = expectedBalances.get(op.account) || 0;
+          expectedBalances.set(op.account, current + op.amount);
+          break;
+        }
+        case 'Withdraw': {
+          api.withdraw(op.account, op.amount);
+          const current = expectedBalances.get(op.account) || 0;
+          expectedBalances.set(op.account, current - op.amount);
+          break;
+        }
+      }
+    }
+
+    // finally check all balances
+    for (const [account, expectedBal] of expectedBalances.entries()) {
+      const retrieved = toNullable(api.balance(account));
+      if (retrieved !== expectedBal) {
+        return false;
+      }
+    }
+
+    return true;
+  };
 
 describe('Bank Simulator', () => {
   test('can retrieve balance given account exists', () => {
@@ -36,6 +112,19 @@ describe('Bank Simulator', () => {
 
     const result = Property.run(prop);
 
+    expect(result.result).toBe('OK');
+  });
+
+  test('correctly update balance given a sequence of deposits and withdrawals', () => {
+    const genTransactionAmount: Gen<number> = choose(1, 100000);
+
+    const prop = forall(genAccounts, accounts => {
+      const api = new SimpleBank(accounts);
+      return property(genOperations(accounts), correctBalance(api, accounts));
+    });
+
+    const result = Property.run(prop);
+
     console.log(result);
     expect(result.result).toBe('OK');
   });