Add tests for disconnect messages

mymarilyn · Aug 26, 2021 · ae4b0ac · ae4b0ac
1 parent 203650e
commit ae4b0ac
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 8 deletions.
diff --git a/clickhouse_driver/connection.py b/clickhouse_driver/connection.py
@@ -479,10 +479,10 @@ def receive_packet(self):
         elif packet_type == ServerPacketTypes.EXCEPTION:
             packet.exception = self.receive_exception()
 
-        elif packet.type == ServerPacketTypes.PROGRESS:
+        elif packet_type == ServerPacketTypes.PROGRESS:
             packet.progress = self.receive_progress()
 
-        elif packet.type == ServerPacketTypes.PROFILE_INFO:
+        elif packet_type == ServerPacketTypes.PROFILE_INFO:
             packet.profile_info = self.receive_profile_info()
 
         elif packet_type == ServerPacketTypes.TOTALS:
@@ -504,12 +504,11 @@ def receive_packet(self):
             )
 
         else:
-            self.disconnect()
-            raise errors.UnknownPacketFromServerError(
-                'Unknown packet {} from server {}'.format(
-                    packet_type, self.get_description()
-                )
+            message = 'Unknown packet {} from server {}'.format(
+                packet_type, self.get_description()
             )
+            self.disconnect()
+            raise errors.UnknownPacketFromServerError(message)
 
         return packet
 

diff --git a/tests/test_connect.py b/tests/test_connect.py
@@ -24,6 +24,12 @@ def test_packets_to_str(self):
 
 
 class ConnectTestCase(BaseTestCase):
+    def unexpected_packet_message(self, expected, got):
+        return (
+            'Code: 102. Unexpected packet from server {}:{} '
+            '(expected {}, got {})'
+        ).format(self.host, self.port, expected, got)
+
     def test_exception_on_hello_packet(self):
         with self.created_client(user='wrong_user') as client:
             with self.assertRaises(errors.ServerException) as e:
@@ -102,9 +108,15 @@ def test_ping_got_unexpected_package(self):
             # Emulate Exception packet on ping.
             mocked_fin.read_one.return_value = 2
 
-            with self.assertRaises(errors.UnexpectedPacketFromServerError):
+            error = errors.UnexpectedPacketFromServerError
+            with self.assertRaises(error) as e:
                 self.client.execute('SELECT 1')
 
+            self.assertEqual(
+                str(e.exception),
+                self.unexpected_packet_message('Pong', 'Exception')
+            )
+
     def test_eof_on_receive_packet(self):
         self.client.execute('SELECT 1')
 
@@ -180,6 +192,36 @@ def test_context_manager(self):
             self.assertTrue(c.connection.connected)
         self.assertFalse(c.connection.connected)
 
+    def test_unknown_packet(self):
+        self.client.execute('SELECT 1')
+
+        with patch('clickhouse_driver.connection.read_varint') as read_mock, \
+                patch.object(self.client.connection, 'force_connect'):
+            read_mock.return_value = 42
+
+            with self.assertRaises(errors.UnknownPacketFromServerError) as e:
+                self.client.execute('SELECT 1')
+
+            self.assertEqual(
+                str(e.exception),
+                'Code: 100. Unknown packet 42 from server {}:{}'.format(
+                    self.host, self.port
+                )
+            )
+
+    def test_unknown_packet_on_connect(self):
+        with patch('clickhouse_driver.connection.read_varint') as read_mock:
+            read_mock.return_value = 42
+
+            error = errors.UnexpectedPacketFromServerError
+            with self.assertRaises(error) as e:
+                self.client.execute('SELECT 1')
+
+            msg = self.unexpected_packet_message(
+                'Hello or Exception', 'Unknown packet'
+            )
+            self.assertEqual(str(e.exception), msg)
+
 
 class FakeBufferedReader(BufferedReader):
     def __init__(self, inputs, bufsize=128):