[ADD] Tests to check if new feature works

Author: josep-tecnativa

tests/test.py

@@ -75,16 +75,30 @@ def _check_local_connection(self):
                         "test_user",
                     ),
                 )
-            except AssertionError:
+                break
+            except ProcessExecutionError:
                 if attempt < 9:
-                    print("Failure number {}. Retrying...".format(attempt))
+                    print(f"Attempt {attempt + 1} failed. Retrying...")
                 else:
                     raise
-            else:
-                continue
+
+    def _execute_sql(self, sql_command):
+        """Execute an SQL command inside the Postgres container."""
+        docker(
+            "container",
+            "exec",
+            self.postgres_container,
+            "psql",
+            "--command",
+            sql_command,
+            "--dbname",
+            "test_db",
+            "--username",
+            "test_user",
+        )
 
     def _check_password_auth(self, host=None):
-        """Test connection with password auth work fine."""
+        """Test connection with password auth works fine."""
         if not host:
             # Connect via LAN by default
             host = self.postgres_container[:12]
@@ -142,7 +156,7 @@ def _check_cert_auth(self):
                 "PGUSER=test_user",
                 CONF_EXTRA,
                 "-v",
-                "{}:/certs".format(local.cwd),
+                f"{local.cwd}:/certs",
                 self.image,
                 "psql",
                 "--host",
@@ -151,6 +165,8 @@ def _check_cert_auth(self):
                 "SELECT 1",
                 "--no-align",
                 "--tuples-only",
+                "--set",
+                "sslmode=verify-full",
             ),
         )
 
@@ -188,7 +204,7 @@ def test_server_certs_mount(self):
             with local.cwd(tdir):
                 self._generate_certs()
                 cert_vols = [
-                    "-v{0}/{1}:/etc/postgres/{1}".format(local.cwd, cert)
+                    f"-v{local.cwd}/{cert}:/etc/postgres/{cert}"
                     for cert in [
                         "client.ca.cert.pem",
                         "server.cert.pem",
@@ -299,6 +315,116 @@ def test_certs_falsy_lan(self):
         with self.assertRaises(ProcessExecutionError):
             self._check_password_auth("example.localdomain")
 
+    def test_hba_extra_rules(self):
+        """Test that HBA_EXTRA_RULES are correctly applied."""
+        # Define custom HBA rules
+        hba_extra_rules = [
+            "host test_db custom_user 0.0.0.0/0 trust",
+            "hostssl all all 192.168.0.0/16 md5",
+        ]
+
+        # Start the Postgres container with HBA_EXTRA_RULES
+        self.postgres_container = docker(
+            "container",
+            "run",
+            "-d",
+            "--network",
+            "lan",
+            "-e",
+            "POSTGRES_DB=test_db",
+            "-e",
+            "POSTGRES_USER=test_user",
+            "-e",
+            "POSTGRES_PASSWORD=test_password",
+            "-e",
+            "HBA_EXTRA_RULES=" + json.dumps(hba_extra_rules),
+            CONF_EXTRA,
+            self.image,
+        ).strip()
+
+        self._check_local_connection()
+
+        # Create custom_user in the database
+        self._execute_sql("CREATE USER custom_user;")
+
+        # Test connection as custom_user without password (trust auth)
+        self.assertEqual(
+            "1\n",
+            docker(
+                "container",
+                "run",
+                "--network",
+                "lan",
+                "-e",
+                "PGDATABASE=test_db",
+                "-e",
+                "PGUSER=custom_user",
+                self.image,
+                "psql",
+                "--host",
+                self.postgres_container[:12],
+                "--command",
+                "SELECT 1",
+                "--no-align",
+                "--tuples-only",
+            ),
+        )
+
+        # Connect the WAN network to test the md5 auth for 192.168.0.0/16
+        self._connect_wan_network(alias="192.168.1.100")
+
+        # Test WAN connection with md5 authentication
+        self._execute_sql("ALTER USER test_user WITH PASSWORD 'test_password';")
+
+        self.assertEqual(
+            "1\n",
+            docker(
+                "container",
+                "run",
+                "--network",
+                "wan",
+                "-e",
+                "PGDATABASE=test_db",
+                "-e",
+                "PGUSER=test_user",
+                "-e",
+                "PGPASSWORD=test_password",
+                "-e",
+                "PGSSLMODE=require",
+                self.image,
+                "psql",
+                "--host",
+                "192.168.1.100",
+                "--command",
+                "SELECT 1",
+                "--no-align",
+                "--tuples-only",
+            ),
+        )
+
+        # Test that connection is refused from WAN with incorrect user
+        with self.assertRaises(ProcessExecutionError):
+            docker(
+                "container",
+                "run",
+                "--network",
+                "wan",
+                "-e",
+                "PGDATABASE=test_db",
+                "-e",
+                "PGUSER=invalid_user",
+                "-e",
+                "PGSSLMODE=require",
+                self.image,
+                "psql",
+                "--host",
+                "192.168.1.100",
+                "--command",
+                "SELECT 1",
+                "--no-align",
+                "--tuples-only",
+            )
+
 
 if __name__ == "__main__":
     unittest.main()