fetchmany ....

TestDataStrategies.py

import unittest, sqlite3
from contextlib import closing

class TestDataStrategies(unittest.TestCase):

    SQLCMD = """SELECT lid
               from person_table
               GROUP BY lid
               ORDER BY lid"""

    def setUp(self):

        self.conn = sqlite3.connect(':memory:')
        self.conn.execute("CREATE TABLE person_table (lid text)")

        patients = [("person_with_id_%s" % x, ) for x in range(9000, 9999)]

        self.conn.executemany("INSERT INTO person_table VALUES (?);", patients)    
        self.conn.commit()

    def test_table_has_data(self):

        with closing(self.conn.cursor()) as cur:
            cur.execute("select * from person_table limit 1")
            pat = cur.fetchone()
            self.assertEqual(len(pat), 1)
            self.assertEqual(pat[0], 'person_with_id_9000')

    def test_batch_retrieval(self):
        with closing(self.conn.cursor()) as cur:
            cur.execute(self.SQLCMD)
            cur.arraysize = 10

            # fetch many without loop returns first set ...
            patients = cur.fetchmany()
            self.assertEqual(len(patients), 10)

            # with `fetchmany`, we have to keep calling fetch many 
            # until no more patients are returned
            patient_counter = 0
            while patients:
                patient_counter += len(patients)
                patients = cur.fetchmany()

            self.assertEqual(patient_counter, 999)


    def tearDown(self):
        self.conn.close()

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