increase coverage

Author: weaverba137

dlairflow/test/test_postgresql.py

@@ -3,8 +3,10 @@
 """Test dlairflow.postgresql.
 """
 import os
+import warnings
 from ..postgresql import pg_dump_schema, pg_restore_schema
 from airflow.hooks.base import BaseHook
+from airflow.operators.bash import BashOperator
 
 
 class MockConnection(object):
@@ -29,7 +31,9 @@ def mock_connection(connection):
 
     test_operator = pg_dump_schema("login,password,host,schema", "dump_schema")
 
-    print(dir(test_operator))
+    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):
@@ -43,4 +47,38 @@ def mock_connection(connection):
 
     test_operator = pg_dump_schema("login,password,host,schema", "dump_schema", "dump_dir")
 
-    print(dir(test_operator))
+    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")
+
+    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):
+    """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'