postgresql: better --host option support

Minor changes: * PEP 3135 "super()" function
jd · Feb 8, 2024 · 8bec8a1 · 8bec8a1
1 parent bcd836e
commit 8bec8a1
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 3 deletions.
diff --git a/pifpaf/drivers/postgresql.py b/pifpaf/drivers/postgresql.py
@@ -40,15 +40,15 @@ def get_options(cls):
     def __init__(self, port=DEFAULT_PORT, host=DEFAULT_HOST,
                  sync=DEFAULT_SYNC, **kwargs):
         """Create a new PostgreSQL instance."""
-        super(PostgreSQLDriver, self).__init__(**kwargs)
+        super().__init__(**kwargs)
         self.port = port
         self.host = host
         self.sync = sync
 
     def _setUp(self):
-        super(PostgreSQLDriver, self)._setUp()
+        super()._setUp()
         self.putenv("PGPORT", str(self.port), True)
-        self.putenv("PGHOST", self.tempdir, True)
+        self.putenv("PGHOST", self.host or self.tempdir, True)
         self.putenv("PGDATA", self.tempdir, True)
         self.putenv("PGDATABASE", "postgres", True)
         _, pgbindir = self._exec(["pg_config", "--bindir"], stdout=True)
@@ -65,6 +65,9 @@ def _setUp(self):
                     % (self.tempdir, self.port, self.host),
                     "start"], allow_debug=False)
         self.addCleanup(self._exec, [pgctl, "-w", "stop"])
+
         self.url = "postgresql://localhost/postgres?host=%s&port=%d" % (
             self.tempdir, self.port)
+        if self.host:
+            self.url = "postgresql://%s:%d/postgres" % (self.host, self.port)
         self.putenv("URL", self.url)
diff --git a/pifpaf/tests/test_drivers.py b/pifpaf/tests/test_drivers.py
@@ -252,6 +252,17 @@ def test_postgresql(self):
             os.getenv("PIFPAF_URL"))
         self._run("psql template1 -c 'CREATE TABLE FOOBAR();'")
 
+    @testtools.skipUnless(spawn.find_executable("pg_config"),
+                          "pg_config not found")
+    def test_postgresql_host(self):
+        host = "127.0.98.25"
+        port = 9825
+        self.useFixture(postgresql.PostgreSQLDriver(host=host, port=port))
+        self.assertEqual(
+            "postgresql://%s:%d/postgres" % (host, port),
+            os.getenv("PIFPAF_URL"))
+        self._run("psql template1 -c 'CREATE TABLE FOOBAR();'")
+
     @testtools.skipUnless(spawn.find_executable("pg_config"),
                           "pg_config not found")
     def test_postgresql_async(self):