diff --git a/dlairflow/test/test_postgresql.py b/dlairflow/test/test_postgresql.py
index c8883e1..149dfe8 100644
--- a/dlairflow/test/test_postgresql.py
+++ b/dlairflow/test/test_postgresql.py
@@ -3,7 +3,7 @@
 """Test dlairflow.postgresql.
 """
 import os
-import warnings
+import pytest
 from ..postgresql import pg_dump_schema, pg_restore_schema
 from airflow.hooks.base import BaseHook
 from airflow.operators.bash import BashOperator
@@ -20,23 +20,9 @@ def __init__(self, connection):
         return
 
 
-def test_pg_dump_schema(monkeypatch):
-    """Test pg_dump task.
-    """
-    def mock_connection(connection):
-        conn = MockConnection(connection)
-        return conn
-
-    monkeypatch.setattr(BaseHook, "get_connection", mock_connection)
-
-    test_operator = pg_dump_schema("login,password,host,schema", "dump_schema")
-
-    assert isinstance(test_operator, BashOperator)
-    assert test_operator.env['PGHOST'] == 'host'
-    assert test_operator.params['schema'] == 'dump_schema'
-
-
-def test_pg_dump_schema_alt_dir(monkeypatch):
+@pytest.mark.parametrize('task_function,dump_dir', [(pg_dump_schema, None), (pg_dump_schema, 'dump_dir'),
+                                                    (pg_restore_schema, None), (pg_restore_schema, 'dump_dir')])
+def test_pg_dump_schema(monkeypatch, task_function, dump_dir):
     """Test pg_dump task with alternate directory.
     """
     def mock_connection(connection):
@@ -45,40 +31,12 @@ def mock_connection(connection):
 
     monkeypatch.setattr(BaseHook, "get_connection", mock_connection)
 
-    test_operator = pg_dump_schema("login,password,host,schema", "dump_schema", "dump_dir")
-
-    assert isinstance(test_operator, BashOperator)
-    assert test_operator.env['PGHOST'] == 'host'
-    assert test_operator.params['dump_dir'] == 'dump_dir'
-
-
-def test_pg_restore_schema(monkeypatch):
-    """Test pg_restore task.
-    """
-    def mock_connection(connection):
-        conn = MockConnection(connection)
-        return conn
-
-    monkeypatch.setattr(BaseHook, "get_connection", mock_connection)
-
-    test_operator = pg_restore_schema("login,password,host,schema", "dump_schema")
+    test_operator = task_function("login,password,host,schema", "dump_schema", dump_dir)
 
     assert isinstance(test_operator, BashOperator)
     assert test_operator.env['PGHOST'] == 'host'
     assert test_operator.params['schema'] == 'dump_schema'
-
-
-def test_pg_restore_schema_alt_dir(monkeypatch):
-    """Test pg_restore task with alternate directory.
-    """
-    def mock_connection(connection):
-        conn = MockConnection(connection)
-        return conn
-
-    monkeypatch.setattr(BaseHook, "get_connection", mock_connection)
-
-    test_operator = pg_restore_schema("login,password,host,schema", "dump_schema", "dump_dir")
-
-    assert isinstance(test_operator, BashOperator)
-    assert test_operator.env['PGHOST'] == 'host'
-    assert test_operator.params['dump_dir'] == 'dump_dir'
+    if dump_dir is None:
+        assert test_operator.params['dump_dir'] == '/data0/datalab/' + os.environ['USER']
+    else:
+        assert test_operator.params['dump_dir'] == 'dump_dir'