From 78c5a89d9462cd424fb267b236e90513d5e495c3 Mon Sep 17 00:00:00 2001
From: Benjamin Alan Weaver <benjamin.weaver@noirlab.edu>
Date: Thu, 12 Dec 2024 17:46:31 -0700
Subject: [PATCH] increase coverage

---
 dlairflow/test/test_postgresql.py | 42 +++++++++++++++++++++++++++++--
 1 file changed, 40 insertions(+), 2 deletions(-)

diff --git a/dlairflow/test/test_postgresql.py b/dlairflow/test/test_postgresql.py
index 1855286..015ed8a 100644
--- a/dlairflow/test/test_postgresql.py
+++ b/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'