Use context manager for socket

... to ensure sockets are closed on unexpected failures.

Change-Id: Iea94fbcfce8e4e2dc9c956eaa3219da0f8173453
This commit is contained in:
Takashi Kajinami 2024-10-06 18:13:01 +09:00
parent 899d78277a
commit 756af007d1
2 changed files with 12 additions and 13 deletions

View File

@ -406,11 +406,9 @@ def get_my_ipv4():
Return ``'127.0.0.1'`` if there is no default interface.
"""
try:
csock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
csock.connect(('192.0.2.0', 80))
(addr, port) = csock.getsockname()
csock.close()
return addr
with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as csock:
csock.connect(('192.0.2.0', 80))
return csock.getsockname()[0]
except socket.error:
return _get_my_ipv4_address()
@ -452,11 +450,9 @@ def get_my_ipv6():
Return ``'::1'`` if there is no default interface.
"""
try:
csock = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
csock.connect(('2001:db8::1', 80))
(addr, _, _, _) = csock.getsockname()
csock.close()
return addr
with socket.socket(socket.AF_INET6, socket.SOCK_DGRAM) as csock:
csock.connect(('2001:db8::1', 80))
return csock.getsockname()[0]
except socket.error:
return _get_my_ipv6_address()

View File

@ -287,16 +287,19 @@ class NetworkUtilsTest(test_base.BaseTestCase):
self.assertFalse(netutils.is_valid_port(input_str))
def test_get_my_ipv4(self):
mock_sock = mock.Mock()
mock_sock.getsockname.return_value = ['1.2.3.4', '']
sock_attrs = {
'return_value.getsockname.return_value': ['1.2.3.4', '']}
'return_value.__enter__.return_value': mock_sock}
with mock.patch('socket.socket', **sock_attrs):
addr = netutils.get_my_ipv4()
self.assertEqual(addr, '1.2.3.4')
def test_get_my_ipv6(self):
mock_sock = mock.Mock()
mock_sock.getsockname.return_value = ['2001:db8::2', '', '', '']
sock_attrs = {
'return_value.getsockname.return_value': ['2001:db8::2', '',
'', '']}
'return_value.__enter__.return_value': mock_sock}
with mock.patch('socket.socket', **sock_attrs):
addr = netutils.get_my_ipv6()
self.assertEqual(addr, '2001:db8::2')