test/python/cursor.py

   1 """
   2 Specialised psycopg2 cursor with shortcut functions useful for testing.
   3 """
   4 import psycopg2.extras
   5
   6 class CursorForTesting(psycopg2.extras.DictCursor):
   7     """ Extension to the DictCursor class that provides execution
   8         short-cuts that simplify writing assertions.
   9     """
  10
  11     def scalar(self, sql, params=None):
  12         """ Execute a query with a single return value and return this value.
  13             Raises an assertion when not exactly one row is returned.
  14         """
  15         self.execute(sql, params)
  16         assert self.rowcount == 1
  17         return self.fetchone()[0]
  18
  19
  20     def row_set(self, sql, params=None):
  21         """ Execute a query and return the result as a set of tuples.
  22             Fails when the SQL command returns duplicate rows.
  23         """
  24         self.execute(sql, params)
  25
  26         result = set((tuple(row) for row in self))
  27         assert len(result) == self.rowcount
  28
  29         return result
  30
  31
  32     def table_exists(self, table):
  33         """ Check that a table with the given name exists in the database.
  34         """
  35         num = self.scalar("""SELECT count(*) FROM pg_tables
  36                              WHERE tablename = %s""", (table, ))
  37         return num == 1
  38
  39
  40     def table_rows(self, table, where=None):
  41         """ Return the number of rows in the given table.
  42         """
  43         if where is None:
  44             return self.scalar('SELECT count(*) FROM ' + table)
  45
  46         return self.scalar('SELECT count(*) FROM {} WHERE {}'.format(table, where))
  47
  48
  49     def execute_values(self, *args, **kwargs):
  50         """ Execute the execute_values() function on the cursor.
  51         """
  52         psycopg2.extras.execute_values(self, *args, **kwargs)