QGIS/python/plugins/db_manager/db_plugins/postgis/plugin_test.py

"""
***************************************************************************
    plugin_test.py
    ---------------------
    Date                 : May 2017
    Copyright            : (C) 2017, Sandro Santilli
    Email                : strk at kbt dot io
***************************************************************************
*                                                                         *
*   This program is free software; you can redistribute it and/or modify  *
*   it under the terms of the GNU General Public License as published by  *
*   the Free Software Foundation; either version 2 of the License, or     *
*   (at your option) any later version.                                   *
*                                                                         *
***************************************************************************
"""

__author__ = "Sandro Santilli"
__date__ = "May 2017"
__copyright__ = "(C) 2017, Sandro Santilli"

import os
import re
import unittest
from qgis.testing import start_app, QgisTestCase
from qgis.core import QgsDataSourceUri
from qgis.utils import iface
from qgis.PyQt.QtCore import QObject

start_app()

from db_manager.db_plugins.postgis.plugin import PostGisDBPlugin, PGRasterTable
from db_manager.db_plugins.postgis.plugin import PGDatabase
from db_manager.db_plugins.postgis.data_model import PGSqlResultModel
from db_manager.db_plugins.plugin import Table
from db_manager.db_plugins.postgis.connector import PostGisDBConnector


class TestDBManagerPostgisPlugin(QgisTestCase):

    @classmethod
    def setUpClass(self):
        self.old_pgdatabase_env = os.environ.get("PGDATABASE")
        # QGIS_PGTEST_DB contains the full connection string and not only the DB name!
        QGIS_PGTEST_DB = os.environ.get("QGIS_PGTEST_DB")
        if QGIS_PGTEST_DB is not None:
            test_uri = QgsDataSourceUri(QGIS_PGTEST_DB)
            self.testdb = test_uri.database()
        else:
            self.testdb = "qgis_test"
        os.environ["PGDATABASE"] = self.testdb

        # Create temporary service file
        self.old_pgservicefile_env = os.environ.get("PGSERVICEFILE")
        self.tmpservicefile = f"/tmp/qgis-test-{os.getpid()}-pg_service.conf"
        os.environ["PGSERVICEFILE"] = self.tmpservicefile

        f = open(self.tmpservicefile, "w")
        f.write(f"[dbmanager]\ndbname={self.testdb}\n")
        # TODO: add more things if PGSERVICEFILE was already set ?
        f.close()

    @classmethod
    def tearDownClass(self):
        # Restore previous env variables if needed
        if self.old_pgdatabase_env:
            os.environ["PGDATABASE"] = self.old_pgdatabase_env
        if self.old_pgservicefile_env:
            os.environ["PGSERVICEFILE"] = self.old_pgservicefile_env
        # Remove temporary service file
        os.unlink(self.tmpservicefile)

    # See https://github.com/qgis/QGIS/issues/24525

    def test_rasterTableURI(self):

        def check_rasterTableURI(expected_dbname):
            tables = database.tables()
            raster_tables_count = 0
            for tab in tables:
                if tab.type == Table.RasterType:
                    raster_tables_count += 1
                    uri = tab.uri()
                    m = re.search(" dbname='([^ ]*)' ", uri)
                    self.assertTrue(m)
                    actual_dbname = m.group(1)
                    self.assertEqual(actual_dbname, expected_dbname)
                # print(tab.type)
                # print(tab.quotedName())
                # print(tab)

            # We need to make sure a database is created with at
            # least one raster table !
            self.assertGreaterEqual(raster_tables_count, 1)

        obj = QObject()  # needs to be kept alive
        obj.connectionName = lambda: "fake"
        obj.providerName = lambda: "postgres"

        # Test for empty URI
        # See https://github.com/qgis/QGIS/issues/24525
        # and https://github.com/qgis/QGIS/issues/19005

        expected_dbname = self.testdb
        os.environ["PGDATABASE"] = expected_dbname

        database = PGDatabase(obj, QgsDataSourceUri())
        self.assertIsInstance(database, PGDatabase)

        uri = database.uri()
        self.assertEqual(uri.host(), "")
        self.assertEqual(uri.username(), "")
        self.assertEqual(uri.database(), expected_dbname)
        self.assertEqual(uri.service(), "")

        check_rasterTableURI(expected_dbname)

        # Test for service-only URI
        # See https://github.com/qgis/QGIS/issues/24526

        os.environ["PGDATABASE"] = "fake"
        database = PGDatabase(obj, QgsDataSourceUri("service=dbmanager"))
        self.assertIsInstance(database, PGDatabase)

        uri = database.uri()
        self.assertEqual(uri.host(), "")
        self.assertEqual(uri.username(), "")
        self.assertEqual(uri.database(), "")
        self.assertEqual(uri.service(), "dbmanager")

        check_rasterTableURI(expected_dbname)

    # See https://github.com/qgis/QGIS/issues/24732
    def test_unicodeInQuery(self):
        os.environ["PGDATABASE"] = self.testdb
        obj = QObject()  # needs to be kept alive
        obj.connectionName = lambda: "fake"
        obj.providerName = lambda: "postgres"
        database = PGDatabase(obj, QgsDataSourceUri())
        self.assertIsInstance(database, PGDatabase)
        # SQL as string literal
        res = database.sqlResultModel("SELECT 'é'::text", obj)
        self.assertIsInstance(res, PGSqlResultModel)
        dat = res.getData(0, 0)
        self.assertEqual(dat, "é")
        # SQL as unicode literal
        res = database.sqlResultModel("SELECT 'é'::text", obj)
        self.assertIsInstance(res, PGSqlResultModel)
        dat = res.getData(0, 0)
        self.assertEqual(dat, "é")


if __name__ == "__main__":
    unittest.main()