adding class for database response

src/py/crankshaft/test/helper.py

@@ -2,6 +2,7 @@ import unittest
 
 from mock_plpy import MockPlPy
 plpy = MockPlPy()
+from mock_plpy import MockDBResponse
 
 import sys
 sys.modules['plpy'] = plpy

src/py/crankshaft/test/mock_plpy.py

@@ -1,12 +1,13 @@
 import re
 
+
 class MockCursor:
     def __init__(self, data):
         self.cursor_pos = 0
         self.data = data
 
     def fetch(self, batch_size):
-        batch = self.data[self.cursor_pos : self.cursor_pos + batch_size]
+        batch = self.data[self.cursor_pos:self.cursor_pos + batch_size]
         self.cursor_pos += batch_size
         return batch
 
@@ -45,8 +46,22 @@ class MockPlPy:
         data = self.execute(query)
         return MockCursor(data)
 
-    def execute(self, query): # TODO: additional arguments
+    # TODO: additional arguments
+    def execute(self, query):
         for result in self.results:
             if result[0].match(query):
                 return result[1]
         return []
+
+
+class MockDBResponse:
+    def __init__(self, data, colnames=None):
+        self.data = data
+        if colnames is None:
+            self.colnames = data[0].keys()
+        else:
+            self.colnames = colnames
+
+
+    def colnames(self):
+        return self.colnames