From b81466d13c30db0d74381913fc7889cb7fc3f7c1 Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Fri, 26 Dec 2025 12:31:43 +0100 Subject: [PATCH v10 3/5] POC: Convert load balance tests from perl to python This is a proof of concept to show how to use the pytest test infrastructure. It converts two existing tests that could not share code. And now they do. If we ever introduce another load balance method (e.g. round robin). We can easily test it for both DNS and hostlist based load balancing by adding a single new test function. --- src/interfaces/libpq/Makefile | 1 + src/interfaces/libpq/meson.build | 7 +- src/interfaces/libpq/pyt/test_load_balance.py | 170 ++++++++++++++++++ .../libpq/t/003_load_balance_host_list.pl | 94 ---------- .../libpq/t/004_load_balance_dns.pl | 144 --------------- 5 files changed, 176 insertions(+), 240 deletions(-) create mode 100644 src/interfaces/libpq/pyt/test_load_balance.py delete mode 100644 src/interfaces/libpq/t/003_load_balance_host_list.pl delete mode 100644 src/interfaces/libpq/t/004_load_balance_dns.pl diff --git a/src/interfaces/libpq/Makefile b/src/interfaces/libpq/Makefile index bf4baa92917..4c4bdb4b3a3 100644 --- a/src/interfaces/libpq/Makefile +++ b/src/interfaces/libpq/Makefile @@ -167,6 +167,7 @@ check installcheck: export PATH := $(CURDIR)/test:$(PATH) check: test-build all $(prove_check) + $(pytest_check) installcheck: test-build all $(prove_installcheck) diff --git a/src/interfaces/libpq/meson.build b/src/interfaces/libpq/meson.build index c5ecd9c3a87..56790dd92a9 100644 --- a/src/interfaces/libpq/meson.build +++ b/src/interfaces/libpq/meson.build @@ -150,8 +150,6 @@ tests += { 'tests': [ 't/001_uri.pl', 't/002_api.pl', - 't/003_load_balance_host_list.pl', - 't/004_load_balance_dns.pl', 't/005_negotiate_encryption.pl', 't/006_service.pl', ], @@ -162,6 +160,11 @@ tests += { }, 'deps': libpq_test_deps, }, + 'pytest': { + 'tests': [ + 'pyt/test_load_balance.py', + ], + }, } subdir('po', if_found: libintl) diff --git a/src/interfaces/libpq/pyt/test_load_balance.py b/src/interfaces/libpq/pyt/test_load_balance.py new file mode 100644 index 00000000000..0af46d8f37d --- /dev/null +++ b/src/interfaces/libpq/pyt/test_load_balance.py @@ -0,0 +1,170 @@ +# Copyright (c) 2025, PostgreSQL Global Development Group + +""" +Tests for load_balance_hosts connection parameter. + +These tests verify that libpq correctly handles load balancing across multiple +PostgreSQL servers specified in the connection string. +""" + +import platform +import re + +import pytest + +from libpq import LibpqError +import pypg + + +@pytest.fixture(scope="module") +def load_balance_nodes_hostlist(create_pg_module): + """ + Create 3 PostgreSQL nodes with different socket directories. + + Each node has its own Unix socket directory for isolation. + Returns a tuple of (nodes, connect). + """ + nodes = [create_pg_module() for _ in range(3)] + + hostlist = ",".join(node.host for node in nodes) + portlist = ",".join(str(node.port) for node in nodes) + + def connect(**kwargs): + return nodes[0].connect(host=hostlist, port=portlist, **kwargs) + + return nodes, connect + + +@pytest.fixture(scope="module") +def load_balance_nodes_dns(create_pg_module): + """ + Create 3 PostgreSQL nodes on the same port but different IP addresses. + + Uses 127.0.0.1, 127.0.0.2, 127.0.0.3 with a shared port, so that + connections to 'pg-loadbalancetest' can be load balanced via DNS. + + Since setting up a DNS server is more effort than we consider reasonable to + run this test, this situation is instead imitated by using a hosts file + where a single hostname maps to multiple different IP addresses. This test + requires the administrator to add the following lines to the hosts file (if + we detect that this hasn't happened we skip the test): + + 127.0.0.1 pg-loadbalancetest + 127.0.0.2 pg-loadbalancetest + 127.0.0.3 pg-loadbalancetest + + Windows or Linux are required to run this test because these OSes allow + binding to 127.0.0.2 and 127.0.0.3 addresses by default, but other OSes + don't. We need to bind to different IP addresses, so that we can use these + different IP addresses in the hosts file. + + The hosts file needs to be prepared before running this test. We don't do + it on the fly, because it requires root permissions to change the hosts + file. In CI we set up the previously mentioned rules in the hosts file, so + that this load balancing method is tested. + + Requires PG_TEST_EXTRA=load_balance because it requires this manual hosts + file configuration and also uses TCP with trust auth, which is potentially + unsafe on multiuser systems. + """ + pypg.skip_unless_test_extras("load_balance") + + if platform.system() not in ("Linux", "Windows"): + pytest.skip("DNS load balance test only supported on Linux and Windows") + + if platform.system() == "Windows": + hosts_path = r"c:\Windows\System32\Drivers\etc\hosts" + else: + hosts_path = "/etc/hosts" + + try: + with open(hosts_path) as f: + hosts_content = f.read() + except (OSError, IOError): + pytest.skip(f"Could not read hosts file: {hosts_path}") + + count = len(re.findall(r"127\.0\.0\.[1-3]\s+pg-loadbalancetest", hosts_content)) + if count != 3: + pytest.skip("hosts file not prepared for DNS load balance test") + + first_node = create_pg_module(hostaddr="127.0.0.1") + nodes = [ + first_node, + create_pg_module(hostaddr="127.0.0.2", port=first_node.port), + create_pg_module(hostaddr="127.0.0.3", port=first_node.port), + ] + + # Allow trust authentication for TCP connections from loopback + for node in nodes: + hba_path = node.datadir / "pg_hba.conf" + with open(hba_path, "r") as f: + original_content = f.read() + with open(hba_path, "w") as f: + f.write("host all all 127.0.0.0/8 trust\n") + f.write(original_content) + node.pg_ctl("reload") + + def connect(**kwargs): + return nodes[0].connect(host="pg-loadbalancetest", **kwargs) + + return nodes, connect + + +@pytest.fixture(scope="module", params=["hostlist", "dns"]) +def load_balance_nodes(request): + """ + Parametrized fixture providing both load balancing test environments. + """ + return request.getfixturevalue(f"load_balance_nodes_{request.param}") + + +def test_load_balance_hosts_invalid_value(load_balance_nodes): + """load_balance_hosts doesn't accept unknown values.""" + _, connect = load_balance_nodes + + with pytest.raises( + LibpqError, match='invalid load_balance_hosts value: "doesnotexist"' + ): + connect(load_balance_hosts="doesnotexist") + + +def test_load_balance_hosts_disable(load_balance_nodes): + """load_balance_hosts=disable always connects to the first node.""" + nodes, connect = load_balance_nodes + + with nodes[0].log_contains("connection received"): + connect(load_balance_hosts="disable") + + +def test_load_balance_hosts_random_distribution(load_balance_nodes): + """load_balance_hosts=random distributes connections across all nodes.""" + nodes, connect = load_balance_nodes + + for _ in range(50): + connect(load_balance_hosts="random") + + occurrences = [ + len(re.findall("connection received", node.log_content())) for node in nodes + ] + + # Statistically, each node should receive at least one connection. + # The probability of any node receiving 0 connections is (2/3)^50 ≈ 1.57e-9 + assert occurrences[0] > 0, "node1 should receive at least one connection" + assert occurrences[1] > 0, "node2 should receive at least one connection" + assert occurrences[2] > 0, "node3 should receive at least one connection" + assert sum(occurrences) == 50, "total connections should be 50" + + +def test_load_balance_hosts_failover(load_balance_nodes): + """load_balance_hosts continues trying hosts until it finds a working one.""" + nodes, connect = load_balance_nodes + + nodes[0].stop() + nodes[1].stop() + + with nodes[2].log_contains("connection received"): + connect(load_balance_hosts="disable") + + with nodes[2].log_contains("connection received", times=5): + for _ in range(5): + connect(load_balance_hosts="random") diff --git a/src/interfaces/libpq/t/003_load_balance_host_list.pl b/src/interfaces/libpq/t/003_load_balance_host_list.pl deleted file mode 100644 index 1f970ff994b..00000000000 --- a/src/interfaces/libpq/t/003_load_balance_host_list.pl +++ /dev/null @@ -1,94 +0,0 @@ -# Copyright (c) 2023-2026, PostgreSQL Global Development Group -use strict; -use warnings FATAL => 'all'; -use Config; -use PostgreSQL::Test::Utils; -use PostgreSQL::Test::Cluster; -use Test::More; - -# This tests load balancing across the list of different hosts in the host -# parameter of the connection string. - -# Cluster setup which is shared for testing both load balancing methods -my $node1 = PostgreSQL::Test::Cluster->new('node1'); -my $node2 = PostgreSQL::Test::Cluster->new('node2', own_host => 1); -my $node3 = PostgreSQL::Test::Cluster->new('node3', own_host => 1); - -# Create a data directory with initdb -$node1->init(); -$node2->init(); -$node3->init(); - -# Start the PostgreSQL server -$node1->start(); -$node2->start(); -$node3->start(); - -# Start the tests for load balancing method 1 -my $hostlist = $node1->host . ',' . $node2->host . ',' . $node3->host; -my $portlist = $node1->port . ',' . $node2->port . ',' . $node3->port; - -$node1->connect_fails( - "host=$hostlist port=$portlist load_balance_hosts=doesnotexist", - "load_balance_hosts doesn't accept unknown values", - expected_stderr => qr/invalid load_balance_hosts value: "doesnotexist"/); - -# load_balance_hosts=disable should always choose the first one. -$node1->connect_ok( - "host=$hostlist port=$portlist load_balance_hosts=disable", - "load_balance_hosts=disable connects to the first node", - sql => "SELECT 'connect1'", - log_like => [qr/statement: SELECT 'connect1'/]); - -# Statistically the following loop with load_balance_hosts=random will almost -# certainly connect at least once to each of the nodes. The chance of that not -# happening is so small that it's negligible: (2/3)^50 = 1.56832855e-9 -foreach my $i (1 .. 50) -{ - $node1->connect_ok( - "host=$hostlist port=$portlist load_balance_hosts=random", - "repeated connections with random load balancing", - sql => "SELECT 'connect2'"); -} - -my $node1_occurrences = () = - $node1->log_content() =~ /statement: SELECT 'connect2'/g; -my $node2_occurrences = () = - $node2->log_content() =~ /statement: SELECT 'connect2'/g; -my $node3_occurrences = () = - $node3->log_content() =~ /statement: SELECT 'connect2'/g; - -my $total_occurrences = - $node1_occurrences + $node2_occurrences + $node3_occurrences; - -cmp_ok($node1_occurrences, '>', 1, - "received at least one connection on node1"); -cmp_ok($node2_occurrences, '>', 1, - "received at least one connection on node2"); -cmp_ok($node3_occurrences, '>', 1, - "received at least one connection on node3"); -is($total_occurrences, 50, "received 50 connections across all nodes"); - -$node1->stop(); -$node2->stop(); - -# load_balance_hosts=disable should continue trying hosts until it finds a -# working one. -$node3->connect_ok( - "host=$hostlist port=$portlist load_balance_hosts=disable", - "load_balance_hosts=disable continues until it connects to the a working node", - sql => "SELECT 'connect3'", - log_like => [qr/statement: SELECT 'connect3'/]); - -# Also with load_balance_hosts=random we continue to the next nodes if previous -# ones are down. Connect a few times to make sure it's not just lucky. -foreach my $i (1 .. 5) -{ - $node3->connect_ok( - "host=$hostlist port=$portlist load_balance_hosts=random", - "load_balance_hosts=random continues until it connects to the a working node", - sql => "SELECT 'connect4'", - log_like => [qr/statement: SELECT 'connect4'/]); -} - -done_testing(); diff --git a/src/interfaces/libpq/t/004_load_balance_dns.pl b/src/interfaces/libpq/t/004_load_balance_dns.pl deleted file mode 100644 index e1ff9a06024..00000000000 --- a/src/interfaces/libpq/t/004_load_balance_dns.pl +++ /dev/null @@ -1,144 +0,0 @@ -# Copyright (c) 2023-2026, PostgreSQL Global Development Group -use strict; -use warnings FATAL => 'all'; -use Config; -use PostgreSQL::Test::Utils; -use PostgreSQL::Test::Cluster; -use Test::More; - -if (!$ENV{PG_TEST_EXTRA} || $ENV{PG_TEST_EXTRA} !~ /\bload_balance\b/) -{ - plan skip_all => - 'Potentially unsafe test load_balance not enabled in PG_TEST_EXTRA'; -} - -# This tests loadbalancing based on a DNS entry that contains multiple records -# for different IPs. Since setting up a DNS server is more effort than we -# consider reasonable to run this test, this situation is instead imitated by -# using a hosts file where a single hostname maps to multiple different IP -# addresses. This test requires the administrator to add the following lines to -# the hosts file (if we detect that this hasn't happened we skip the test): -# -# 127.0.0.1 pg-loadbalancetest -# 127.0.0.2 pg-loadbalancetest -# 127.0.0.3 pg-loadbalancetest -# -# Windows or Linux are required to run this test because these OSes allow -# binding to 127.0.0.2 and 127.0.0.3 addresses by default, but other OSes -# don't. We need to bind to different IP addresses, so that we can use these -# different IP addresses in the hosts file. -# -# The hosts file needs to be prepared before running this test. We don't do it -# on the fly, because it requires root permissions to change the hosts file. In -# CI we set up the previously mentioned rules in the hosts file, so that this -# load balancing method is tested. - -# Cluster setup which is shared for testing both load balancing methods -my $can_bind_to_127_0_0_2 = - $Config{osname} eq 'linux' || $PostgreSQL::Test::Utils::windows_os; - -# Checks for the requirements for testing load balancing method 2 -if (!$can_bind_to_127_0_0_2) -{ - plan skip_all => 'load_balance test only supported on Linux and Windows'; -} - -my $hosts_path; -if ($windows_os) -{ - $hosts_path = 'c:\Windows\System32\Drivers\etc\hosts'; -} -else -{ - $hosts_path = '/etc/hosts'; -} - -my $hosts_content = PostgreSQL::Test::Utils::slurp_file($hosts_path); - -my $hosts_count = () = - $hosts_content =~ /127\.0\.0\.[1-3] pg-loadbalancetest/g; -if ($hosts_count != 3) -{ - # Host file is not prepared for this test - plan skip_all => "hosts file was not prepared for DNS load balance test"; -} - -$PostgreSQL::Test::Cluster::use_tcp = 1; -$PostgreSQL::Test::Cluster::test_pghost = '127.0.0.1'; -my $port = PostgreSQL::Test::Cluster::get_free_port(); -my $node1 = PostgreSQL::Test::Cluster->new('node1', port => $port); -my $node2 = - PostgreSQL::Test::Cluster->new('node2', port => $port, own_host => 1); -my $node3 = - PostgreSQL::Test::Cluster->new('node3', port => $port, own_host => 1); - -# Create a data directory with initdb -$node1->init(); -$node2->init(); -$node3->init(); - -# Start the PostgreSQL server -$node1->start(); -$node2->start(); -$node3->start(); - -# load_balance_hosts=disable should always choose the first one. -$node1->connect_ok( - "host=pg-loadbalancetest port=$port load_balance_hosts=disable", - "load_balance_hosts=disable connects to the first node", - sql => "SELECT 'connect1'", - log_like => [qr/statement: SELECT 'connect1'/]); - - -# Statistically the following loop with load_balance_hosts=random will almost -# certainly connect at least once to each of the nodes. The chance of that not -# happening is so small that it's negligible: (2/3)^50 = 1.56832855e-9 -foreach my $i (1 .. 50) -{ - $node1->connect_ok( - "host=pg-loadbalancetest port=$port load_balance_hosts=random", - "repeated connections with random load balancing", - sql => "SELECT 'connect2'"); -} - -my $node1_occurrences = () = - $node1->log_content() =~ /statement: SELECT 'connect2'/g; -my $node2_occurrences = () = - $node2->log_content() =~ /statement: SELECT 'connect2'/g; -my $node3_occurrences = () = - $node3->log_content() =~ /statement: SELECT 'connect2'/g; - -my $total_occurrences = - $node1_occurrences + $node2_occurrences + $node3_occurrences; - -cmp_ok($node1_occurrences, '>', 1, - "received at least one connection on node1"); -cmp_ok($node2_occurrences, '>', 1, - "received at least one connection on node2"); -cmp_ok($node3_occurrences, '>', 1, - "received at least one connection on node3"); -is($total_occurrences, 50, "received 50 connections across all nodes"); - -$node1->stop(); -$node2->stop(); - -# load_balance_hosts=disable should continue trying hosts until it finds a -# working one. -$node3->connect_ok( - "host=pg-loadbalancetest port=$port load_balance_hosts=disable", - "load_balance_hosts=disable continues until it connects to a working node", - sql => "SELECT 'connect3'", - log_like => [qr/statement: SELECT 'connect3'/]); - -# Also with load_balance_hosts=random we continue to the next nodes if previous -# ones are down. Connect a few times to make sure it's not just lucky. -foreach my $i (1 .. 5) -{ - $node3->connect_ok( - "host=pg-loadbalancetest port=$port load_balance_hosts=random", - "load_balance_hosts=random continues until it connects to a working node", - sql => "SELECT 'connect4'", - log_like => [qr/statement: SELECT 'connect4'/]); -} - -done_testing(); -- 2.52.0