From 6c8b66df96146c5b1e5c30e337d9b02fb0aa337d Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sun, 21 Sep 2025 09:52:39 -0600 Subject: [PATCH] cleanup --- tests/unit_tests/test_espota2.py | 64 ++++++++++++++++---------------- 1 file changed, 33 insertions(+), 31 deletions(-) diff --git a/tests/unit_tests/test_espota2.py b/tests/unit_tests/test_espota2.py index 80fb804bc8..fd46ae3b49 100644 --- a/tests/unit_tests/test_espota2.py +++ b/tests/unit_tests/test_espota2.py @@ -14,9 +14,20 @@ from esphome import espota2 from esphome.core import EsphomeError -def test_recv_decode_with_decode() -> None: +@pytest.fixture +def mock_socket(): + """Create a mock socket for testing.""" + socket = Mock() + socket.close = Mock() + socket.recv = Mock() + socket.sendall = Mock() + socket.settimeout = Mock() + socket.connect = Mock() + return socket + + +def test_recv_decode_with_decode(mock_socket) -> None: """Test recv_decode with decode=True returns list.""" - mock_socket = Mock() mock_socket.recv.return_value = b"\x01\x02\x03" result = espota2.recv_decode(mock_socket, 3, decode=True) @@ -25,9 +36,8 @@ def test_recv_decode_with_decode() -> None: mock_socket.recv.assert_called_once_with(3) -def test_recv_decode_without_decode() -> None: +def test_recv_decode_without_decode(mock_socket) -> None: """Test recv_decode with decode=False returns bytes.""" - mock_socket = Mock() mock_socket.recv.return_value = b"\x01\x02\x03" result = espota2.recv_decode(mock_socket, 3, decode=False) @@ -36,9 +46,8 @@ def test_recv_decode_without_decode() -> None: mock_socket.recv.assert_called_once_with(3) -def test_receive_exactly_success() -> None: +def test_receive_exactly_success(mock_socket) -> None: """Test receive_exactly successfully receives expected data.""" - mock_socket = Mock() mock_socket.recv.side_effect = [b"\x00", b"\x01\x02"] result = espota2.receive_exactly(mock_socket, 3, "test", espota2.RESPONSE_OK) @@ -47,11 +56,9 @@ def test_receive_exactly_success() -> None: assert mock_socket.recv.call_count == 2 -def test_receive_exactly_with_error_response() -> None: +def test_receive_exactly_with_error_response(mock_socket) -> None: """Test receive_exactly raises OTAError on error response.""" - mock_socket = Mock() mock_socket.recv.return_value = bytes([espota2.RESPONSE_ERROR_AUTH_INVALID]) - mock_socket.close = Mock() with pytest.raises(espota2.OTAError, match="Error auth:.*Authentication invalid"): espota2.receive_exactly(mock_socket, 1, "auth", [espota2.RESPONSE_OK]) @@ -59,9 +66,8 @@ def test_receive_exactly_with_error_response() -> None: mock_socket.close.assert_called_once() -def test_receive_exactly_socket_error() -> None: +def test_receive_exactly_socket_error(mock_socket) -> None: """Test receive_exactly handles socket errors.""" - mock_socket = Mock() mock_socket.recv.side_effect = OSError("Connection reset") with pytest.raises(espota2.OTAError, match="Error receiving acknowledge test"): @@ -119,9 +125,8 @@ def test_check_error_unexpected_response() -> None: espota2.check_error([0x7F], [espota2.RESPONSE_OK, espota2.RESPONSE_AUTH_OK]) -def test_send_check_with_various_data_types() -> None: +def test_send_check_with_various_data_types(mock_socket) -> None: """Test send_check handles different data types.""" - mock_socket = Mock() # Test with list/tuple espota2.send_check(mock_socket, [0x01, 0x02], "list") @@ -140,18 +145,16 @@ def test_send_check_with_various_data_types() -> None: mock_socket.sendall.assert_called_with(b"\xaa\xbb") -def test_send_check_socket_error() -> None: +def test_send_check_socket_error(mock_socket) -> None: """Test send_check handles socket errors.""" - mock_socket = Mock() mock_socket.sendall.side_effect = OSError("Broken pipe") with pytest.raises(espota2.OTAError, match="Error sending test"): espota2.send_check(mock_socket, b"data", "test") -def test_perform_ota_successful_md5_auth() -> None: +def test_perform_ota_successful_md5_auth(mock_socket) -> None: """Test successful OTA with MD5 authentication.""" - mock_socket = Mock() mock_file = io.BytesIO(b"firmware content here") # Mock random for predictable cnonce @@ -183,9 +186,14 @@ def test_perform_ota_successful_md5_auth() -> None: # Verify magic bytes were sent assert mock_socket.sendall.call_args_list[0] == call(bytes(espota2.MAGIC_BYTES)) - # Verify features were sent + # Verify features were sent (compression + SHA256 support) assert mock_socket.sendall.call_args_list[1] == call( - bytes([espota2.FEATURE_SUPPORTS_COMPRESSION]) + bytes( + [ + espota2.FEATURE_SUPPORTS_COMPRESSION + | espota2.FEATURE_SUPPORTS_SHA256_AUTH + ] + ) ) # Verify cnonce was sent (MD5 of random.random()) @@ -201,9 +209,8 @@ def test_perform_ota_successful_md5_auth() -> None: assert mock_socket.sendall.call_args_list[3] == call(expected_result.encode()) -def test_perform_ota_no_auth() -> None: +def test_perform_ota_no_auth(mock_socket) -> None: """Test OTA without authentication.""" - mock_socket = Mock() mock_file = io.BytesIO(b"firmware") with patch("time.sleep"), patch("time.perf_counter", side_effect=[0, 1]): @@ -231,9 +238,8 @@ def test_perform_ota_no_auth() -> None: assert len(auth_calls) == 0 -def test_perform_ota_with_compression() -> None: +def test_perform_ota_with_compression(mock_socket) -> None: """Test OTA with compression support.""" - mock_socket = Mock() original_content = b"firmware" * 100 # Repeating content for compression mock_file = io.BytesIO(original_content) @@ -274,9 +280,8 @@ def test_perform_ota_with_compression() -> None: assert sent_size == len(compressed) -def test_perform_ota_auth_without_password() -> None: +def test_perform_ota_auth_without_password(mock_socket) -> None: """Test OTA fails when auth is required but no password provided.""" - mock_socket = Mock() mock_file = io.BytesIO(b"firmware") responses = [ @@ -293,9 +298,8 @@ def test_perform_ota_auth_without_password() -> None: espota2.perform_ota(mock_socket, "", mock_file, "test.bin") -def test_perform_ota_unsupported_version() -> None: +def test_perform_ota_unsupported_version(mock_socket) -> None: """Test OTA fails with unsupported version.""" - mock_socket = Mock() mock_file = io.BytesIO(b"firmware") responses = [ @@ -308,9 +312,8 @@ def test_perform_ota_unsupported_version() -> None: espota2.perform_ota(mock_socket, "", mock_file, "test.bin") -def test_perform_ota_upload_error() -> None: +def test_perform_ota_upload_error(mock_socket) -> None: """Test OTA handles upload errors.""" - mock_socket = Mock() mock_file = io.BytesIO(b"firmware") with patch("time.perf_counter", side_effect=[0, 1]): @@ -511,9 +514,8 @@ def test_perform_ota_sha256_fallback_to_md5() -> None: pass # Implementation depends on final PR merge -def test_perform_ota_version_differences() -> None: +def test_perform_ota_version_differences(mock_socket) -> None: """Test OTA behavior differences between version 1.0 and 2.0.""" - mock_socket = Mock() mock_file = io.BytesIO(b"firmware") with patch("time.sleep"), patch("time.perf_counter", side_effect=[0, 1]):