import sys

try:
    from pymysql.tests import base
    import pymysql.cursors
    from pymysql.constants import CLIENT
except Exception:
    # For local testing from top-level directory, without installing
    sys.path.append('../pymysql')
    from pymysql.tests import base
    import pymysql.cursors
    from pymysql.constants import CLIENT

class TestSSCursor(base.PyMySQLTestCase):
    def test_SSCursor(self):
        affected_rows = 18446744073709551615

        conn = self.connect(client_flag=CLIENT.MULTI_STATEMENTS)
        data = [
            ('America', '', 'America/Jamaica'),
            ('America', '', 'America/Los_Angeles'),
            ('America', '', 'America/Lima'),
            ('America', '', 'America/New_York'),
            ('America', '', 'America/Menominee'),
            ('America', '', 'America/Havana'),
            ('America', '', 'America/El_Salvador'),
            ('America', '', 'America/Costa_Rica'),
            ('America', '', 'America/Denver'),
            ('America', '', 'America/Detroit'),]

        cursor = conn.cursor(pymysql.cursors.SSCursor)

        # Create table
        cursor.execute('CREATE TABLE tz_data ('
            'region VARCHAR(64),'
            'zone VARCHAR(64),'
            'name VARCHAR(64))')

        conn.begin()
        # Test INSERT
        for i in data:
            cursor.execute('INSERT INTO tz_data VALUES (%s, %s, %s)', i)
            self.assertEqual(conn.affected_rows(), 1, 'affected_rows does not match')
        conn.commit()

        # Test fetchone()
        iter = 0
        cursor.execute('SELECT * FROM tz_data')
        while True:
            row = cursor.fetchone()
            if row is None:
                break
            iter += 1

            # Test cursor.rowcount
            self.assertEqual(cursor.rowcount, affected_rows,
                'cursor.rowcount != %s' % (str(affected_rows)))

            # Test cursor.rownumber
            self.assertEqual(cursor.rownumber, iter,
                'cursor.rowcount != %s' % (str(iter)))

            # Test row came out the same as it went in
            self.assertEqual((row in data), True,
                'Row not found in source data')

        # Test fetchall
        cursor.execute('SELECT * FROM tz_data')
        self.assertEqual(len(cursor.fetchall()), len(data),
            'fetchall failed. Number of rows does not match')

        # Test fetchmany
        cursor.execute('SELECT * FROM tz_data')
        self.assertEqual(len(cursor.fetchmany(2)), 2,
            'fetchmany failed. Number of rows does not match')

        # So MySQLdb won't throw "Commands out of sync"
        while True:
            res = cursor.fetchone()
            if res is None:
                break

        # Test update, affected_rows()
        cursor.execute('UPDATE tz_data SET zone = %s', ['Foo'])
        conn.commit()
        self.assertEqual(cursor.rowcount, len(data),
            'Update failed. affected_rows != %s' % (str(len(data))))

        # Test executemany
        cursor.executemany('INSERT INTO tz_data VALUES (%s, %s, %s)', data)
        self.assertEqual(cursor.rowcount, len(data),
            'executemany failed. cursor.rowcount != %s' % (str(len(data))))

        # Test multiple datasets
        cursor.execute('SELECT 1; SELECT 2; SELECT 3')
        self.assertListEqual(list(cursor), [(1, )])
        self.assertTrue(cursor.nextset())
        self.assertListEqual(list(cursor), [(2, )])
        self.assertTrue(cursor.nextset())
        self.assertListEqual(list(cursor), [(3, )])
        self.assertFalse(cursor.nextset())

        cursor.execute('DROP TABLE IF EXISTS tz_data')
        cursor.close()

__all__ = ["TestSSCursor"]

if __name__ == "__main__":
    import unittest
    unittest.main()