Skip to content

Commit de1ce5c

Browse files
committed
Merge branch 'fix-initdb-error' of github.com:postgrespro/testgres into fix-initdb-error
2 parents 036b924 + 05cd996 commit de1ce5c

File tree

5 files changed

+68
-23
lines changed

5 files changed

+68
-23
lines changed

testgres/helpers/__init__.py

Whitespace-only changes.

testgres/helpers/port_manager.py

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
import socket
2+
import random
3+
from typing import Set, Iterable, Optional
4+
5+
6+
class PortForException(Exception):
7+
pass
8+
9+
10+
class PortManager:
11+
def __init__(self, ports_range=(1024, 65535)):
12+
self.ports_range = ports_range
13+
14+
@staticmethod
15+
def is_port_free(port: int) -> bool:
16+
"""Check if a port is free to use."""
17+
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
18+
try:
19+
s.bind(("", port))
20+
return True
21+
except OSError:
22+
return False
23+
24+
def find_free_port(self, ports: Optional[Set[int]] = None, exclude_ports: Optional[Iterable[int]] = None) -> int:
25+
"""Return a random unused port number."""
26+
if ports is None:
27+
ports = set(range(1024, 65535))
28+
29+
if exclude_ports is None:
30+
exclude_ports = set()
31+
32+
ports.difference_update(set(exclude_ports))
33+
34+
sampled_ports = random.sample(tuple(ports), min(len(ports), 100))
35+
36+
for port in sampled_ports:
37+
if self.is_port_free(port):
38+
return port
39+
40+
raise PortForException("Can't select a port")

testgres/operations/local_ops.py

Lines changed: 25 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -39,23 +39,30 @@ def __init__(self, conn_params=None):
3939
self.username = conn_params.username or self.get_user()
4040

4141
@staticmethod
42-
def _run_command(cmd, shell, input, timeout, encoding, temp_file=None):
42+
def _run_command(cmd, shell, input, stdin, stdout, stderr, timeout, encoding, temp_file=None, get_process=None):
4343
"""Execute a command and return the process."""
4444
if temp_file is not None:
45-
stdout = temp_file
46-
stderr = subprocess.STDOUT
45+
stdout = stdout or temp_file
46+
stderr = stderr or subprocess.STDOUT
4747
else:
48-
stdout = subprocess.PIPE
49-
stderr = subprocess.PIPE
48+
stdout = stdout or subprocess.PIPE
49+
stderr = stderr or subprocess.PIPE
5050

5151
process = subprocess.Popen(
5252
cmd,
5353
shell=shell,
54-
stdin=subprocess.PIPE if input is not None else None,
54+
stdin=stdin or subprocess.PIPE if input is not None else None,
5555
stdout=stdout,
5656
stderr=stderr,
5757
)
58-
return process
58+
59+
if get_process:
60+
return None, process
61+
try:
62+
return process.communicate(input=input.encode(encoding) if input else None, timeout=timeout), process
63+
except subprocess.TimeoutExpired:
64+
process.kill()
65+
raise ExecUtilException("Command timed out after {} seconds.".format(timeout))
5966

6067
@staticmethod
6168
def _raise_exec_exception(message, command, exit_code, output):
@@ -67,7 +74,7 @@ def _raise_exec_exception(message, command, exit_code, output):
6774

6875
def exec_command(self, cmd, wait_exit=False, verbose=False,
6976
expect_error=False, encoding=None, shell=False, text=False,
70-
input=None, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
77+
input=None, stdin=None, stdout=None, stderr=None,
7178
get_process=None, timeout=None):
7279
"""
7380
Execute a command in a subprocess.
@@ -143,29 +150,27 @@ def _process_output(process, encoding, temp_file=None):
143150

144151
def _exec_command_windows(self, cmd, wait_exit=False, verbose=False,
145152
expect_error=False, encoding=None, shell=False, text=False,
146-
input=None, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
153+
input=None, stdin=None, stdout=None, stderr=None,
147154
get_process=None, timeout=None):
148155
with tempfile.NamedTemporaryFile(mode='w+b') as temp_file:
149-
process = self._run_command(cmd, shell, input, timeout, encoding, temp_file)
150-
try:
151-
output = process.communicate(input=input.encode(encoding) if input else None, timeout=timeout)
152-
except subprocess.TimeoutExpired:
153-
process.kill()
154-
raise ExecUtilException("Command timed out after {} seconds.".format(timeout))
156+
_, process = self._run_command(cmd, shell, input, stdin, stdout, stderr, timeout, encoding, temp_file, get_process)
157+
if get_process:
158+
return process
159+
result = self._process_output(process, encoding, temp_file)
155160

156-
if process.returncode != 0 or has_errors(output):
161+
if process.returncode != 0 or has_errors(result):
157162
if process.returncode == 0:
158163
process.returncode = 1
159164
if expect_error:
160165
if verbose:
161-
return process.returncode, output, output
166+
return process.returncode, result, result
162167
else:
163-
return output
168+
return result
164169
else:
165170
self._raise_exec_exception('Utility exited with non-zero code. Error `{}`', cmd, process.returncode,
166-
output)
171+
result)
167172

168-
return (process.returncode, output, output) if verbose else output
173+
return (process.returncode, result, result) if verbose else result
169174

170175
# Environment setup
171176
def environ(self, var_name):

testgres/utils.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313

1414
from six import iteritems
1515

16-
from helpers.port_manager import PortManager
16+
from .helpers.port_manager import PortManager
1717
from .exceptions import ExecUtilException
1818
from .config import testgres_config as tconf
1919

tests/test_simple.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,7 @@ def test_init_after_cleanup(self):
141141
node.cleanup()
142142
node.init().start().execute('select 1')
143143

144-
@unittest.skipUnless(util_exists('pg_resetwal'), 'might be missing')
144+
@unittest.skipUnless(util_exists('pg_resetwal.exe' if os.name == 'nt' else 'pg_resetwal'), 'pgbench might be missing')
145145
@unittest.skipUnless(pg_version_ge('9.6'), 'requires 9.6+')
146146
def test_init_unique_system_id(self):
147147
# this function exists in PostgreSQL 9.6+
@@ -747,7 +747,7 @@ def test_logging(self):
747747
master.restart()
748748
self.assertTrue(master._logger.is_alive())
749749

750-
@unittest.skipUnless(util_exists('pgbench'), 'might be missing')
750+
@unittest.skipUnless(util_exists('pgbench.exe' if os.name == 'nt' else 'pgbench'), 'pgbench might be missing')
751751
def test_pgbench(self):
752752
with get_new_node().init().start() as node:
753753

0 commit comments

Comments
 (0)