more unit tests

chimpler · Apr 29, 2015 · 0c1c074 · 0c1c074
1 parent 8b5a340
commit 0c1c074
Show file tree

Hide file tree

Showing 6 changed files with 67 additions and 14 deletions.
diff --git a/README.md b/README.md
@@ -206,8 +206,8 @@ Data export                            | :white_check_mark:
 Data import                            | :white_check_mark:
 Constraints export                     | :x:
 Constraints import                     | :x:
-Postgres                               | :x:
-MySQL                                  | :x:
+Postgres                               | :white_check_mark:
+MySQL                                  | :white_check_mark:
 Oracle                                 | :x:
 Vertica                                | :x:
 SQLite                                 | :x:

diff --git a/catdb/main.py b/catdb/main.py
@@ -3,7 +3,6 @@
 import os
 from pyhocon import ConfigFactory
 import sys
-from catdb import CatDbException
 from catdb.db import DbManager
 
 

diff --git a/catdb/mysql.py b/catdb/mysql.py
@@ -45,10 +45,6 @@ def get_col_def(row):
                 'scale': scale
             }
 
-        conn = self.get_connection()
-        try:
-            with conn.cursor() as cursor:
-                cursor.execute('DESC ' + table)
-                return [get_col_def(row) for row in cursor.fetchall()]
-        finally:
-            conn.close()
+        with self.get_connection(False) as cursor:
+            cursor.execute('DESC ' + table)
+            return [get_col_def(row) for row in cursor.fetchall()]
diff --git a/catdb/postgres.py b/catdb/postgres.py
@@ -5,10 +5,6 @@
 class Postgres(Db):
     PUBLIC_SCHEMA = 'public'
 
-    DEFAULT_VALUE_MAPPING = {
-        'CURRENT_TIMESTAMP': 'NOW()'
-    }
-
     def __init__(self, params):
         super(Postgres, self).__init__('postgres', params)
 

diff --git a/tests/test_mysql.py b/tests/test_mysql.py
@@ -20,3 +20,32 @@ def test_list_tables(self, mock_pymysql):
             'table_a',
             'table_b'
         ]
+
+    @mock.patch('catdb.mysql.pymysql')
+    def test_get_column_info(self, mock_pymysql):
+        mock_cursor = mock.MagicMock()
+        mock_cursor.fetchall.return_value = [
+            {
+                'Field': 'field',
+                'Type': 'VARCHAR(254)',
+                'Null': 'YES',
+                'Default': 'test'
+            }
+        ]
+        mock_pymysql.connect.return_value.__enter__.return_value = mock_cursor
+
+        mysql = Mysql({
+            'database': 'test'
+        })
+
+        assert mysql.get_column_info(None, 'test') == [
+            {
+                'column': 'field',
+                'type': 'varchar',
+                'size': 254,
+                'radix': None,
+                'nullable': True,
+                'default': 'test',
+                'scale': None,
+            }
+        ]
diff --git a/tests/test_postgres.py b/tests/test_postgres.py
@@ -21,3 +21,36 @@ def test_list_tables(self, mock_psycopg2):
             'table_a',
             'table_b'
         ]
+
+    @mock.patch('catdb.postgres.psycopg2')
+    def test_get_column_info(self, mock_psycopg2):
+        mock_cursor = mock.MagicMock()
+        mock_cursor.cursor.return_value.__enter__.return_value.fetchall.return_value = [
+            [
+                'field',  # column
+                'character varying',  # data_type
+                'test',  # default_value
+                True,  # nullable
+                254,  # length
+                None,  # numeric_precision
+                None,  # numeric_precision_radix
+                None  # numeric_scale
+            ]
+        ]
+        mock_psycopg2.connect.return_value.__enter__.return_value = mock_cursor
+
+        postgres = Postgres({
+            'database': 'test'
+        })
+
+        assert postgres.get_column_info(None, 'test') == [
+            {
+                'column': 'field',
+                'type': 'character varying',
+                'size': 254,
+                'radix': None,
+                'nullable': True,
+                'default': 'test',
+                'scale': None,
+            }
+        ]