Exercise: Mocking A Bank


exdb/bank.py
import db

class Bank():
    def __init__(self):
        self.db = db.DB()

    def setup(self):
        self.db.create()

    def transfer(self, src, dst, amount):
        src_current = self.db.get(src)
        dst_current = self.db.get(dst)
        if src_current and src_current >= amount:
            self.db.update(src, src_current-amount)
            self.db.update(dst, dst_current+amount)
        else:
            raise Exception("Not enough money")

    def status(self, name):
        return self.db.get(name)
  
    def deposit(self, name, amount):
        current = self.db.get(name)
        if current == None:
            self.db.insert(name, amount)
        else:
            self.db.update(name, current+amount)


exdb/db.py
import sqlite3

db_filename = 'bank.db'


class DB():
    def __init__(self):
        self.db_filename = db_filename
        self.conn = sqlite3.connect(self.db_filename)

    def create(self):
        c = self.conn.cursor()
        c.execute('''CREATE TABLE account
                 (name text, ballance real)''')

    def get(self, name):
        c = self.conn.cursor()
        c.execute('SELECT ballance FROM account WHERE name=?', (name,))
        current = c.fetchone()
        if current == None:
            return current
        else:
            return current[0]

    def insert(self, name, amount):
        c = self.conn.cursor()
        c.execute('INSERT INTO account (name, ballance) VALUES (?, ?)', (name, amount))

    def update(self, name, amount):
        c = self.conn.cursor()
        c.execute('UPDATE account SET ballance = ? WHERE name = ?', (amount, name))