sqlite_test.py

#!/usr/bin/env python

import sqlite3, unittest

class KittehDB:
    def __init__(self,path):
        """
        Creates database at specified path
        """
        print("Creating Database")
        self.conn   = sqlite3.connect(path)
        self.cursor = self.conn.cursor()
        self.__execute("""
        CREATE TABLE IF NOT EXISTS kittehs(name TEXT, weight INTEGER)
        """)
                
    def __del__(self):
        """
        Close DB Connection (for :memory: this means freeing data)
        """
        print("Closing database.")
        self.conn.close()

    def __execute(self,command,*args):
        """
        Works same as:
        self.cursor.execute() but does an additional commit
        """
        self.cursor.execute(command,*args)
        self.conn.commit()

    def insert(self,name = 'myKitteh',weight = 0):
        """
        Insert a value to the DB
        """
        self.__execute('INSERT INTO kittehs VALUES (?,?)',(name,weight))

    def get_contents(self):
        """ 
        Simply queries all data and writes set
        """
        s = {}
        self.cursor.execute('SELECT * FROM kittehs')
        for row in self.cursor:
            s[row[0]] = row[1]

        return s


class KittehDBTest(unittest.TestCase):
    def setUp(self):
        self.db = KittehDB(':memory:')


    def test_insert(self):
        num = 10000
        for i in range(num):
            self.db.insert(str(i),i)

        kDict = self.db.get_contents()
        self.assertEqual(kDict.__len__(),num)

        for i in range(num):
            self.assertEqual(kDict[str(i)],i)




if __name__ == '__main__':
    unittest.main()