From 0c1c07465577436a5f04db9d749cdc93922a76b7 Mon Sep 17 00:00:00 2001
From: Francois Dang Ngoc <francois.dangngoc@gmail.com>
Date: Tue, 28 Apr 2015 23:14:30 -0400
Subject: [PATCH] more unit tests

---
 README.md              |  4 ++--
 catdb/main.py          |  1 -
 catdb/mysql.py         | 10 +++-------
 catdb/postgres.py      |  4 ----
 tests/test_mysql.py    | 29 +++++++++++++++++++++++++++++
 tests/test_postgres.py | 33 +++++++++++++++++++++++++++++++++
 6 files changed, 67 insertions(+), 14 deletions(-)

diff --git a/README.md b/README.md
index b19dbdb..e823658 100644
--- 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
index 660fbfd..7e6c5b4 100644
--- 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
index 0343e28..eaf2909 100644
--- 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
index 4934709..315d706 100644
--- 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
index f38708c..16e8624 100644
--- 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
index 20d9298..c1c51f7 100644
--- 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,
+            }
+        ]