#!/usr/bin/env python

import logging
import unittest
import apypg
import time
import traceback

class SimpleTestCase(unittest.TestCase):
    def setUp(self):
        self.conn = apypg.Connection()

    def testEmptyQuery(self):
        cursor = self.conn.query("select 'x' where false")
        for row in cursor.tuples():
            self.fail("Shouldn't have returned any rows!")

    def testAlmostEmptyQuery(self):
        cursor = self.conn.query("select 'x'")
        self.assertEquals([['x']], list(cursor.tuples()))

    def testParseError(self):
        try:
            cursor = self.conn.query("foo")
            self.assertEquals([], list(cursor.tuples()))
        except Exception, e:
            self.failIf(str(e).find('syntax error') == -1)
        else:
            self.fail("Really should have errored")

    def testMultipleQueries(self):
        cursor = self.conn.query("select 'x'")
        self.assertEquals([['x']], list(cursor.tuples()))
        cursor = self.conn.query("select 'y'")
        self.assertEquals([['y']], list(cursor.tuples()))

    def testLargeQuery(self):
        """query large enough to see incremental results"""
        cursor = self.conn.query("select count(*) from sale_component")
        count = int(list(cursor.tuples())[0][0])
        self.failUnless(count > 1000, "Need to load sales data")
        start = time.time()
        cursor = self.conn.query("select * from sale_component")
        i = 0
        firstrow = None
        for _ in cursor.tuples():
            if firstrow is None:
                firstrow = time.time()
            i += 1
        lastrow = time.time()
        self.assertEquals(count, i)
        self.assert_(lastrow - start > 100*(firstrow-start),
            "first row too slow! firstrow=%f lastrow=%f"
                % (firstrow-start, lastrow-start))

    def testShortQueryError(self):
        # The error row will probably be in the first packet of rows, so
        # I expect we'll not process any rows before getting the error.
        self._testQueryError(5,10)

    def testLongQueryError(self):
        # We'll certainly get a chance to peek at rows before seeing this
        # error.
        self._testQueryError(10000,20000)

    def _testQueryError(self, error_row, num_rows):
        # http://archives.postgresql.org/pgsql-interfaces/2005-11/msg00015.php
        cursor = self.conn.query("""
            select    1/(generate_series-%(error_row)d)
            from      generate_series(1,%(num_rows)d)
        """ % {'error_row': error_row, 'num_rows': num_rows})
        i = 0
        try:
            for row in cursor.tuples():
                i += 1
        except Exception, e:
            self.failIf(str(e).find('division by zero') == -1,
                'Unexpected error %s' % e)
            # XXX: I think the best plan is for this to always give all
            # rows before the error. For now, I just check that it doesn't
            # give extra rows. It'd require a libpq change to have the
            # more strict behavior.
            self.assert_(i < error_row, 'Errored _after_ row %d' % error_row)
        else:
            self.fail('Should have caused error')

if __name__ == '__main__':
    logging.basicConfig()
    logging.getLogger().setLevel(logging.DEBUG)
    unittest.main()