From 4778b275d56a954c39370b86dfde76f644301cd5 Mon Sep 17 00:00:00 2001 From: Kai Chappell Date: Thu, 19 Jun 2025 15:55:38 +0000 Subject: [PATCH] Add transport layer tests --- src/py_dvt_ate/instruments/transport/tcp.py | 2 + tests/unit/test_transport.py | 263 ++++++++++++++++++++ 2 files changed, 265 insertions(+) create mode 100644 tests/unit/test_transport.py diff --git a/src/py_dvt_ate/instruments/transport/tcp.py b/src/py_dvt_ate/instruments/transport/tcp.py index 7076298..32b82f0 100644 --- a/src/py_dvt_ate/instruments/transport/tcp.py +++ b/src/py_dvt_ate/instruments/transport/tcp.py @@ -147,6 +147,8 @@ class TCPTransport: # Decode and strip whitespace return response_bytes.decode(self._encoding).strip() + except ConnectionError: + raise # Re-raise ConnectionError as-is except socket.timeout as err: raise TimeoutError("Read timeout") from err except (socket.error, OSError, UnicodeDecodeError) as err: diff --git a/tests/unit/test_transport.py b/tests/unit/test_transport.py new file mode 100644 index 0000000..e372ea6 --- /dev/null +++ b/tests/unit/test_transport.py @@ -0,0 +1,263 @@ +"""Unit tests for transport layer.""" + +import socket +from unittest.mock import MagicMock, Mock, patch + +import pytest + +from py_dvt_ate.instruments.transport.tcp import TCPTransport + + +class TestTCPTransport: + """Tests for TCPTransport class.""" + + @pytest.fixture + def transport(self) -> TCPTransport: + """Create transport instance for tests.""" + return TCPTransport("localhost", 5025, timeout=1.0) + + def test_creation(self, transport: TCPTransport) -> None: + """Test TCPTransport can be created with valid parameters.""" + assert transport.host == "localhost" + assert transport.port == 5025 + assert not transport.is_connected + + def test_repr(self, transport: TCPTransport) -> None: + """Test string representation.""" + assert "localhost:5025" in repr(transport) + assert "disconnected" in repr(transport) + + @patch("socket.socket") + def test_connect_success( + self, mock_socket_class: Mock, transport: TCPTransport + ) -> None: + """Test successful connection.""" + mock_sock = MagicMock() + mock_socket_class.return_value = mock_sock + + transport.connect() + + assert transport.is_connected + mock_socket_class.assert_called_once_with( + socket.AF_INET, socket.SOCK_STREAM + ) + mock_sock.settimeout.assert_called_once_with(1.0) + mock_sock.connect.assert_called_once_with(("localhost", 5025)) + + @patch("socket.socket") + def test_connect_already_connected( + self, mock_socket_class: Mock, transport: TCPTransport + ) -> None: + """Test connecting when already connected raises error.""" + mock_sock = MagicMock() + mock_socket_class.return_value = mock_sock + + transport.connect() + + with pytest.raises(ConnectionError, match="Already connected"): + transport.connect() + + @patch("socket.socket") + def test_connect_failure( + self, mock_socket_class: Mock, transport: TCPTransport + ) -> None: + """Test connection failure raises ConnectionError.""" + mock_sock = MagicMock() + mock_socket_class.return_value = mock_sock + mock_sock.connect.side_effect = socket.error("Connection refused") + + with pytest.raises(ConnectionError, match="Failed to connect"): + transport.connect() + + assert not transport.is_connected + + @patch("socket.socket") + def test_disconnect( + self, mock_socket_class: Mock, transport: TCPTransport + ) -> None: + """Test disconnect closes socket.""" + mock_sock = MagicMock() + mock_socket_class.return_value = mock_sock + + transport.connect() + transport.disconnect() + + mock_sock.close.assert_called_once() + assert not transport.is_connected + + def test_disconnect_when_not_connected(self, transport: TCPTransport) -> None: + """Test disconnect is idempotent.""" + transport.disconnect() # Should not raise + assert not transport.is_connected + + @patch("socket.socket") + def test_write(self, mock_socket_class: Mock, transport: TCPTransport) -> None: + """Test write sends command with newline.""" + mock_sock = MagicMock() + mock_socket_class.return_value = mock_sock + + transport.connect() + transport.write("*IDN?") + + mock_sock.sendall.assert_called_once_with(b"*IDN?\n") + + def test_write_not_connected(self, transport: TCPTransport) -> None: + """Test write when not connected raises error.""" + with pytest.raises(ConnectionError, match="Not connected"): + transport.write("*IDN?") + + @patch("socket.socket") + def test_write_failure( + self, mock_socket_class: Mock, transport: TCPTransport + ) -> None: + """Test write failure raises IOError.""" + mock_sock = MagicMock() + mock_socket_class.return_value = mock_sock + mock_sock.sendall.side_effect = socket.error("Write failed") + + transport.connect() + + with pytest.raises(OSError, match="Write failed"): + transport.write("*IDN?") + + @patch("socket.socket") + def test_read(self, mock_socket_class: Mock, transport: TCPTransport) -> None: + """Test read receives response until newline.""" + mock_sock = MagicMock() + mock_socket_class.return_value = mock_sock + + # Simulate receiving "OK\n" byte by byte + mock_sock.recv.side_effect = [b"O", b"K", b"\n"] + + transport.connect() + response = transport.read() + + assert response == "OK" + + @patch("socket.socket") + def test_read_with_timeout( + self, mock_socket_class: Mock, transport: TCPTransport + ) -> None: + """Test read with custom timeout.""" + mock_sock = MagicMock() + mock_socket_class.return_value = mock_sock + mock_sock.gettimeout.return_value = 1.0 + + # Simulate receiving "OK\n" + mock_sock.recv.side_effect = [b"O", b"K", b"\n"] + + transport.connect() + response = transport.read(timeout=2.0) + + assert response == "OK" + # Verify timeout was changed and restored + assert mock_sock.settimeout.call_count == 3 # connect + custom + restore + mock_sock.settimeout.assert_any_call(2.0) + mock_sock.settimeout.assert_any_call(1.0) + + def test_read_not_connected(self, transport: TCPTransport) -> None: + """Test read when not connected raises error.""" + with pytest.raises(ConnectionError, match="Not connected"): + transport.read() + + @patch("socket.socket") + def test_read_timeout( + self, mock_socket_class: Mock, transport: TCPTransport + ) -> None: + """Test read timeout raises TimeoutError.""" + mock_sock = MagicMock() + mock_socket_class.return_value = mock_sock + mock_sock.recv.side_effect = socket.timeout("Timed out") + + transport.connect() + + with pytest.raises(TimeoutError, match="Read timeout"): + transport.read() + + @patch("socket.socket") + def test_read_connection_closed( + self, mock_socket_class: Mock, transport: TCPTransport + ) -> None: + """Test read when connection closed raises error.""" + mock_sock = MagicMock() + mock_socket_class.return_value = mock_sock + mock_sock.recv.return_value = b"" # Empty means connection closed + + transport.connect() + + with pytest.raises(ConnectionError, match="Connection closed"): + transport.read() + + @patch("socket.socket") + def test_query(self, mock_socket_class: Mock, transport: TCPTransport) -> None: + """Test query combines write and read.""" + mock_sock = MagicMock() + mock_socket_class.return_value = mock_sock + + # Simulate receiving "Test Device\n" + mock_sock.recv.side_effect = [ + b"T", + b"e", + b"s", + b"t", + b" ", + b"D", + b"e", + b"v", + b"i", + b"c", + b"e", + b"\n", + ] + + transport.connect() + response = transport.query("*IDN?") + + assert response == "Test Device" + mock_sock.sendall.assert_called_once_with(b"*IDN?\n") + + @patch("socket.socket") + def test_query_with_timeout( + self, mock_socket_class: Mock, transport: TCPTransport + ) -> None: + """Test query with custom timeout.""" + mock_sock = MagicMock() + mock_socket_class.return_value = mock_sock + mock_sock.gettimeout.return_value = 1.0 + + mock_sock.recv.side_effect = [b"O", b"K", b"\n"] + + transport.connect() + response = transport.query("*IDN?", timeout=3.0) + + assert response == "OK" + mock_sock.settimeout.assert_any_call(3.0) + + @patch("socket.socket") + def test_context_manager( + self, mock_socket_class: Mock, transport: TCPTransport + ) -> None: + """Test context manager connects and disconnects.""" + mock_sock = MagicMock() + mock_socket_class.return_value = mock_sock + + with transport: + assert transport.is_connected + + mock_sock.close.assert_called_once() + assert not transport.is_connected + + @patch("socket.socket") + def test_context_manager_with_exception( + self, mock_socket_class: Mock, transport: TCPTransport + ) -> None: + """Test context manager disconnects even on exception.""" + mock_sock = MagicMock() + mock_socket_class.return_value = mock_sock + + with pytest.raises(ValueError): + with transport: + raise ValueError("Test error") + + mock_sock.close.assert_called_once() + assert not transport.is_connected