From 4afbbef034f46c769a302d370ede89fc566ba5f2 Mon Sep 17 00:00:00 2001
From: Fantix King <fantix.king@gmail.com>
Date: Tue, 20 Feb 2024 21:20:11 -0500
Subject: [PATCH] Add test

---
 tests/test_protocol.py | 36 +++++++++++++++++++++++++++++-------
 1 file changed, 29 insertions(+), 7 deletions(-)

diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index 1635288fc2c..15e8ef8fe60 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -404,14 +404,14 @@ async def test_proto_desc_id_cardinality(self):
 
         self.assertNotEqual(cdd1.output_typedesc_id, cdd2.output_typedesc_id)
 
-    async def _parse(self, query):
+    async def _parse(self, query, output_format=protocol.OutputFormat.BINARY):
         await self.con.send(
             protocol.Parse(
                 annotations=[],
                 allowed_capabilities=protocol.Capability.ALL,
                 compilation_flags=protocol.CompilationFlag(0),
                 implicit_limit=0,
-                output_format=protocol.OutputFormat.BINARY,
+                output_format=output_format,
                 expected_cardinality=compiler.Cardinality.MANY,
                 command_text=query,
                 state_typedesc_id=b'\0' * 16,
@@ -511,9 +511,8 @@ async def test_proto_state_concurrent_alter(self):
             await con2.aclose()
 
     async def _parse_execute(self, query, args):
-        await self.con.connect()
-
-        await self._parse(query)
+        output_format=protocol.OutputFormat.BINARY
+        await self._parse(query, output_format=output_format)
         res = await self.con.recv()
 
         await self.con.send(
@@ -523,7 +522,7 @@ async def _parse_execute(self, query, args):
                 compilation_flags=protocol.CompilationFlag(0),
                 implicit_limit=0,
                 command_text=query,
-                output_format=protocol.OutputFormat.NONE,
+                output_format=output_format,
                 expected_cardinality=protocol.Cardinality.MANY,
                 input_typedesc_id=res.input_typedesc_id,
                 output_typedesc_id=res.output_typedesc_id,
@@ -533,7 +532,6 @@ async def _parse_execute(self, query, args):
             ),
             protocol.Sync(),
         )
-        await self.con.recv()
 
     async def test_proto_execute_bad_array_01(self):
         q = "SELECT <array<int32>>$0"
@@ -558,6 +556,7 @@ async def test_proto_execute_bad_array_01(self):
             len(array),   # len
         ) + array
 
+        await self.con.connect()
         await self._parse_execute(q, args)
         await self.con.recv_match(
             protocol.ErrorResponse,
@@ -586,6 +585,7 @@ async def test_proto_execute_bad_array_02(self):
             len(array),   # len
         ) + array
 
+        await self.con.connect()
         await self._parse_execute(q, args)
         await self.con.recv_match(
             protocol.ErrorResponse,
@@ -614,6 +614,7 @@ async def test_proto_execute_bad_array_03(self):
             len(array),   # len
         ) + array
 
+        await self.con.connect()
         await self._parse_execute(q, args)
         await self.con.recv_match(
             protocol.ErrorResponse,
@@ -657,6 +658,27 @@ async def test_proto_global_bad_array(self):
             message='invalid NULL'
         )
 
+    async def test_proto_parse_execute_transaction_id(self):
+        await self.con.connect()
+        await self._parse_execute("start transaction", b"")
+        await self.con.recv_match(
+            protocol.CommandComplete,
+            status='START TRANSACTION'
+        )
+        await self.con.recv_match(
+            protocol.ReadyForCommand,
+            transaction_state=protocol.TransactionState.IN_TRANSACTION,
+        )
+        await self._parse_execute("commit", b"")
+        await self.con.recv_match(
+            protocol.CommandComplete,
+            status='COMMIT'
+        )
+        await self.con.recv_match(
+            protocol.ReadyForCommand,
+            transaction_state=protocol.TransactionState.NOT_IN_TRANSACTION,
+        )
+
 
 class TestServerCancellation(tb.TestCase):
     @contextlib.asynccontextmanager