Fix tests for CopyCommand / UnloadCommand

Author: Robert Engel

tests/test_copy_command.py

@@ -1,18 +1,18 @@
 from unittest import TestCase
-from sqlalchemy import Table, Column, Integer, String, MetaData
-from sqlalchemy.sql import select, func
 from redshift_sqlalchemy.dialect import CopyCommand
 
 
 class TestCopyCommand(TestCase):
     def setUp(self):
-        ''' Sets up a table and associate meta data for the test queries to build against
-        '''
         pass
 
     def test_basic_copy_case(self):
-        ''' Tests that the simplest type of UnloadFromSelect works
+        ''' Tests that the simplest type of CopyCommand works
         '''
-        expected_result = "copy t1 from 's3://mybucket/data/listing/' credentials 'aws_access_key_id=cookies;aws_secret_access_key=cookies'"
-        insert = CopyCommand('t1', 's3://mybucket/data/listing/', 'cookies', 'cookies')
-        self.assertEqual(expected_result, str(insert))
+        expected_result = "COPY schema1.t1 FROM 's3://mybucket/data/listing/'\n" \
+                          "           CREDENTIALS 'aws_access_key_id=cookies;aws_secret_access_key=cookies'\n" \
+                          "           CSV\n           TRUNCATECOLUMNS\n           DELIMITER ','\n" \
+                          "           IGNOREHEADER 0\n           NULL '---'\n           EMPTYASNULL\n" \
+                          "           BLANKSASNULL;"
+        insert = CopyCommand('schema1', 't1', 's3://mybucket/data/listing/', 'cookies', 'cookies')
+        self.assertEqual(expected_result, str(insert).strip())

tests/test_unload_from_select.py

@@ -14,13 +14,21 @@ def setUp(self):
     def test_basic_unload_case(self):
         ''' Tests that the simplest type of UnloadFromSelect works
         '''
-        expected_result = "unload ('SELECT count(t1.id) AS count_1 \nFROM t1') to 'cookies' credentials 'aws_access_key_id=cookies;aws_secret_access_key=cookies' delimiter ',' addquotes allowoverwrite parallel on"
-        insert = UnloadFromSelect(select([func.count(self.t1.c.id)]), 'cookies', 'cookies', 'cookies')
-        self.assertEqual(expected_result, str(insert))
+        expected_result = "UNLOAD ('SELECT count(t1.id) AS count_1 \nFROM t1') TO 's3://bucket/key'\n           " \
+                          "CREDENTIALS 'aws_access_key_id=cookies;aws_secret_access_key=cookies;token=cookies'\n" \
+                          "           DELIMITER ','\n           ADDQUOTES\n           ALLOWOVERWRITE\n          " \
+                          " PARALLEL ON;"
+        insert = UnloadFromSelect(select([func.count(self.t1.c.id)]), 's3://bucket/key', 'cookies', 'cookies',
+                                  'cookies')
+        self.assertEqual(expected_result, str(insert).strip())
 
     def test_parallel_off_unload_case(self):
         ''' Tests that UnloadFromSelect handles parallel being set to off
         '''
-        expected_result = "unload ('SELECT count(t1.id) AS count_1 \nFROM t1') to 'cookies' credentials 'aws_access_key_id=cookies;aws_secret_access_key=cookies' delimiter ',' addquotes allowoverwrite parallel off"
-        insert = UnloadFromSelect(select([func.count(self.t1.c.id)]), 'cookies', 'cookies', 'cookies', 'off')
-        self.assertEqual(expected_result, str(insert))
+        expected_result = "UNLOAD ('SELECT count(t1.id) AS count_1 \nFROM t1') TO 's3://bucket/key'\n           " \
+                          "CREDENTIALS 'aws_access_key_id=cookies;aws_secret_access_key=cookies;token=cookies'\n" \
+                          "           DELIMITER ','\n           ADDQUOTES\n           ALLOWOVERWRITE\n          " \
+                          " PARALLEL OFF;"
+        insert = UnloadFromSelect(select([func.count(self.t1.c.id)]), 's3://bucket/key', 'cookies', 'cookies',
+                                  'cookies', {'parallel': 'OFF'})
+        self.assertEqual(expected_result, str(insert).strip())