module_utils.postgres: added unittests (#56381)

lib/ansible/module_utils/postgres.py

@@ -27,9 +27,9 @@
 # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 # USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+psycopg2 = None  # This line needs for unit tests
 try:
     import psycopg2
-    from psycopg2.extras import DictCursor
     HAS_PSYCOPG2 = True
 except ImportError:
     HAS_PSYCOPG2 = False
@@ -55,6 +55,11 @@ def ensure_libs(sslrootcert=None):
 
 
 def postgres_common_argument_spec():
+    """
+    Return a dictionary with connection options.
+
+    The options are commonly used by most of PostgreSQL modules.
+    """
     return dict(
         login_user=dict(default='postgres'),
         login_password=dict(default='', no_log=True),
@@ -67,6 +72,7 @@ def postgres_common_argument_spec():
 
 
 def ensure_required_libs(module):
+    """Check required libraries."""
     if not HAS_PSYCOPG2:
         module.fail_json(msg=missing_required_lib('psycopg2'))
 
@@ -75,7 +81,14 @@ def ensure_required_libs(module):
 
 
 def connect_to_db(module, autocommit=False, fail_on_conn=True, warn_db_default=True):
+    """Return psycopg2 connection object.
 
+    Keyword arguments:
+    module -- object of ansible.module_utils.basic.AnsibleModule class
+    autocommit -- commit automatically (default False)
+    fail_on_conn -- fail if connection failed or just warn and return None (default True)
+    warn_db_default -- warn that the default DB is used (default True)
+    """
     ensure_required_libs(module)
 
     # To use defaults values, keyword arguments must be absent, so
@@ -110,22 +123,24 @@ def connect_to_db(module, autocommit=False, fail_on_conn=True, warn_db_default=T
     if is_localhost and module.params["login_unix_socket"] != "":
         kw["host"] = module.params["login_unix_socket"]
 
+    db_connection = None
     try:
         db_connection = psycopg2.connect(**kw)
         if autocommit:
-            if psycopg2.__version__ >= '2.4.2':
+            if LooseVersion(psycopg2.__version__) >= LooseVersion('2.4.2'):
                 db_connection.set_session(autocommit=True)
             else:
                 db_connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
 
         # Switch role, if specified:
-        cursor = db_connection.cursor(cursor_factory=DictCursor)
         if module.params.get('session_role'):
+            cursor = db_connection.cursor(cursor_factory=psycopg2.extras.DictCursor)
             try:
                 cursor.execute('SET ROLE %s' % module.params['session_role'])
             except Exception as e:
                 module.fail_json(msg="Could not switch role: %s" % to_native(e))
-        cursor.close()
+            finally:
+                cursor.close()
 
     except TypeError as e:
         if 'sslrootcert' in e.args[0]: