3 files changed, 126 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');
   });
diff --git a/pbt/ts/src/bank.ts b/pbt/ts/src/bank.ts
index c5adc05..e5619f4 100644
--- a/pbt/ts/src/bank.ts
+++ b/pbt/ts/src/bank.ts
@@ -39,3 +39,9 @@ export class SimpleBank implements Bank {
     throw new Error('Method not implemented.');
   }
 }
+
+export type Operation = GetBalance | Withdraw | Deposit;
+
+export type GetBalance = {op: 'GetBalance'; account: Account};
+export type Withdraw = {op: 'Withdraw'; account: Account; amount: number};
+export type Deposit = {op: 'Deposit'; account: Account; amount: number};
diff --git a/pbt/ts/src/property.ts b/pbt/ts/src/property.ts
index 500c423..04f2a62 100644
--- a/pbt/ts/src/property.ts
+++ b/pbt/ts/src/property.ts
@@ -1,3 +1,5 @@
+import {Option, getOrElse, isSome, none, some} from 'fp-ts/lib/Option';
+import {LazyArg, pipe} from 'fp-ts/lib/function';
 import Prando from 'prando';
 
 type Result = 'OK' | 'Falsified' | 'Exception';
@@ -125,6 +127,31 @@ export function elements<A>(arr: A[]): Gen<A> {
   };
 }
 
+export function frequency<A>(distribution: [number, Gen<A>][]): Gen<A> {
+  const total: number = distribution.reduce((acc, [freq, _]) => acc + freq, 0);
+
+  return (rng: Prando) => (size: number) => {
+    const idx = rng.next(0, total);
+    let selectedGen: Option<Gen<A>> = none;
+    selectedGen = distribution.reduce((found, [freq, gen]) => {
+      if (isSome(found)) {
+        return found;
+      } else if (idx <= freq) {
+        return some(gen);
+      } else {
+        return none;
+      }
+    }, selectedGen);
+    const generator: Gen<A> = pipe(
+      selectedGen,
+      getOrElse((() => {
+        throw new Error(`invalid frequency distribution ${distribution}`);
+      }) as LazyArg<Gen<A>>),
+    );
+    return generator(rng)(size);
+  };
+}
+
 /// Select one key-value pair from given map
 export function keyValueOf<K, V>(map: Map<K, V>): Gen<[K, V]> {
   let keys = Array.from(map.keys());