Fix tests expecting to warn without server_hostname

2 files changed, 8 insertions, 4 deletions

diff --git a/test/test_ssl.py b/test/test_ssl.py
index 8634d7bd..82171234 100644
--- a/test/test_ssl.py
+++ b/test/test_ssl.py
@@ -62,9 +62,12 @@ def test_sni_missing_warning_with_ip_addresses(monkeypatch, has_sni, server_host
     sock = mock.Mock()
     context = mock.Mock()
 
-    with pytest.warns(None) as record:
+    with mock.patch('warnings.warn') as warn:
         ssl_.ssl_wrap_socket(sock, server_hostname=server_hostname, ssl_context=context)
 
-    assert len(record) == int(should_warn)
     if should_warn:
-        assert isinstance(record[0].message, SNIMissingWarning)
+        assert warn.call_count >= 1
+        warnings = [call[0][1] for call in warn.call_args_list]
+        assert SNIMissingWarning in warnings
+    else:
+        assert warn.call_count == 0
diff --git a/test/test_util.py b/test/test_util.py
index c2e2fef2..73d9452a 100644
--- a/test/test_util.py
+++ b/test/test_util.py
@@ -500,7 +500,8 @@ class TestUtil(object):
         ssl_.HAS_SNI = False
         try:
             with patch('warnings.warn') as warn:
-                ssl_wrap_socket(ssl_context=mock_context, sock=socket)
+                ssl_wrap_socket(ssl_context=mock_context, sock=socket,
+                                server_hostname='www.google.com')
             mock_context.wrap_socket.assert_called_once_with(socket)
             assert warn.call_count >= 1
             warnings = [call[0][1] for call in warn.call_args_list]