Merge pull request googleapis#3414 from GoogleCloudPlatform/3019-spanner-streaming_chunking_systests

tseaver · web-flow · commit df6700871023 · 2017-05-16T10:34:00.000-04:00
Test streaming of various payload row sizes.
diff --git a/spanner/tests/system/test_system.py b/spanner/tests/system/test_system.py
@@ -726,6 +726,59 @@ def test_execute_sql_w_query_param(self):
         )
 
 
+class TestStreamingChunking(unittest.TestCase, _TestData):
+
+    @classmethod
+    def setUpClass(cls):
+        from tests.system.utils.streaming_utils import INSTANCE_NAME
+        from tests.system.utils.streaming_utils import DATABASE_NAME
+
+        instance = Config.CLIENT.instance(INSTANCE_NAME)
+        if not instance.exists():
+            raise unittest.SkipTest(
+                "Run 'tests/system/utils/populate_streaming.py' to enable.")
+
+        database = instance.database(DATABASE_NAME)
+        if not instance.exists():
+            raise unittest.SkipTest(
+                "Run 'tests/system/utils/populate_streaming.py' to enable.")
+
+        cls._db = database
+
+    def _verify_one_column(self, table_desc):
+        sql = 'SELECT chunk_me FROM {}'.format(table_desc.table)
+        rows = list(self._db.execute_sql(sql))
+        self.assertEqual(len(rows), table_desc.row_count)
+        expected = table_desc.value()
+        for row in rows:
+            self.assertEqual(row[0], expected)
+
+    def _verify_two_columns(self, table_desc):
+        sql = 'SELECT chunk_me, chunk_me_2 FROM {}'.format(table_desc.table)
+        rows = list(self._db.execute_sql(sql))
+        self.assertEqual(len(rows), table_desc.row_count)
+        expected = table_desc.value()
+        for row in rows:
+            self.assertEqual(row[0], expected)
+            self.assertEqual(row[1], expected)
+
+    def test_four_kay(self):
+        from tests.system.utils.streaming_utils import FOUR_KAY
+        self._verify_one_column(FOUR_KAY)
+
+    def test_forty_kay(self):
+        from tests.system.utils.streaming_utils import FOUR_KAY
+        self._verify_one_column(FOUR_KAY)
+
+    def test_four_hundred_kay(self):
+        from tests.system.utils.streaming_utils import FOUR_HUNDRED_KAY
+        self._verify_one_column(FOUR_HUNDRED_KAY)
+
+    def test_four_meg(self):
+        from tests.system.utils.streaming_utils import FOUR_MEG
+        self._verify_two_columns(FOUR_MEG)
+
+
 class _DatabaseDropper(object):
     """Helper for cleaning up databases created on-the-fly."""
 
diff --git a/spanner/tests/system/utils/populate_streaming.py b/spanner/tests/system/utils/populate_streaming.py
@@ -19,29 +19,35 @@
 from google.cloud.spanner.pool import BurstyPool
 
 # Import relative to the script's directory
+from streaming_utils import FOUR_KAY
+from streaming_utils import FORTY_KAY
+from streaming_utils import FOUR_HUNDRED_KAY
+from streaming_utils import FOUR_MEG
 from streaming_utils import DATABASE_NAME
 from streaming_utils import INSTANCE_NAME
 from streaming_utils import print_func
 
+
 DDL = """\
-CREATE TABLE four_kay (
+CREATE TABLE {0.table} (
     pkey INT64,
-    chunk_me STRING(4096) )
+    chunk_me STRING({0.value_size}) )
     PRIMARY KEY (pkey);
-CREATE TABLE forty_kay (
+CREATE TABLE {1.table} (
     pkey INT64,
-    chunk_me STRING(40960) )
+    chunk_me STRING({1.value_size}) )
     PRIMARY KEY (pkey);
-CREATE TABLE four_hundred_kay (
+CREATE TABLE {2.table} (
     pkey INT64,
-    chunk_me STRING(409600) )
+    chunk_me STRING({2.value_size}) )
     PRIMARY KEY (pkey);
-CREATE TABLE four_meg (
+CREATE TABLE {3.table} (
     pkey INT64,
-    chunk_me STRING(2097152),
-    chunk_me_2 STRING(2097152) )
+    chunk_me STRING({3.value_size}),
+    chunk_me_2 STRING({3.value_size}) )
     PRIMARY KEY (pkey);
-"""
+""".format(FOUR_KAY, FORTY_KAY, FOUR_HUNDRED_KAY, FOUR_MEG)
+
 
 DDL_STATEMENTS = [stmt.strip() for stmt in DDL.split(';') if stmt.strip()]
 
@@ -75,49 +81,51 @@ def ensure_database(client):
     return database
 
 
-def populate_table(database, table_name, row_count, val_size):
+def populate_table(database, table_desc):
     all_ = KeySet(all_=True)
     columns = ('pkey', 'chunk_me')
     rows = list(database.execute_sql(
-        'SELECT COUNT(*) FROM {}'.format(table_name)))
+        'SELECT COUNT(*) FROM {}'.format(table_desc.table)))
     assert len(rows) == 1
     count = rows[0][0]
-    if count != row_count:
-        print_func("Repopulating table: {}".format(table_name))
-        chunk_me = 'X' * val_size
-        row_data = [(index, chunk_me) for index in range(row_count)]
+    if count != table_desc.row_count:
+        print_func("Repopulating table: {}".format(table_desc.table))
+        chunk_me = table_desc.value()
+        row_data = [(index, chunk_me) for index in range(table_desc.row_count)]
         with database.batch() as batch:
-            batch.delete(table_name, all_)
-            batch.insert(table_name, columns, row_data)
+            batch.delete(table_desc.table, all_)
+            batch.insert(table_desc.table, columns, row_data)
     else:
-        print_func("Leaving table: {}".format(table_name))
+        print_func("Leaving table: {}".format(table_desc.table))
 
 
-def populate_table_2_columns(database, table_name, row_count, val_size):
+def populate_table_2_columns(database, table_desc):
     all_ = KeySet(all_=True)
     columns = ('pkey', 'chunk_me', 'chunk_me_2')
     rows = list(database.execute_sql(
-        'SELECT COUNT(*) FROM {}'.format(table_name)))
+        'SELECT COUNT(*) FROM {}'.format(table_desc.table)))
     assert len(rows) == 1
     count = rows[0][0]
-    if count != row_count:
-        print_func("Repopulating table: {}".format(table_name))
-        chunk_me = 'X' * val_size
-        row_data = [(index, chunk_me, chunk_me) for index in range(row_count)]
+    if count != table_desc.row_count:
+        print_func("Repopulating table: {}".format(table_desc.table))
+        chunk_me = table_desc.value()
+        row_data = [
+            (index, chunk_me, chunk_me)
+            for index in range(table_desc.row_count)]
         with database.batch() as batch:
-            batch.delete(table_name, all_)
-            batch.insert(table_name, columns, row_data)
+            batch.delete(table_desc.table, all_)
+            batch.insert(table_desc.table, columns, row_data)
     else:
-        print_func("Leaving table: {}".format(table_name))
+        print_func("Leaving table: {}".format(table_desc.table))
 
 
 def populate_streaming(client):
     database = ensure_database(client)
-    populate_table(database, 'four_kay', 1000, 4096)
-    populate_table(database, 'forty_kay', 100, 4096 * 10)
-    populate_table(database, 'four_hundred_kay', 25, 4096 * 100)
+    populate_table(database, FOUR_KAY)
+    populate_table(database, FORTY_KAY)
+    populate_table(database, FOUR_HUNDRED_KAY)
     # Max STRING column size is just larger than 2 Mb, so use two columns
-    populate_table_2_columns(database, 'four_meg', 10, 2048 * 1024)
+    populate_table_2_columns(database, FOUR_MEG)
 
 
 if __name__ == '__main__':
diff --git a/spanner/tests/system/utils/streaming_utils.py b/spanner/tests/system/utils/streaming_utils.py
@@ -12,13 +12,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import collections
 import os
 
 INSTANCE_NAME = 'gcp-streaming-systests'
 DATABASE_NAME = 'testing'
 _SHOULD_PRINT = os.getenv('GOOGLE_CLOUD_NO_PRINT') != 'true'
 
 
+class _TableDesc(collections.namedtuple(
+    'TableDesc', ('table', 'row_count', 'value_size', 'column_count'))):
+
+    def value(self):
+        return u'X' * self.value_size
+
+
+FOUR_KAY = _TableDesc('four_kay', 1000, 4096, 1)
+FORTY_KAY = _TableDesc('forty_kay', 100, 4096 * 10, 1)
+FOUR_HUNDRED_KAY = _TableDesc('four_hundred_kay', 25, 4096 * 100, 1)
+FOUR_MEG = _TableDesc('four_meg', 10, 2048 *  1024, 2)
+
+
 def print_func(message):
     if _SHOULD_PRINT:
         print(message)
