tx context manager tests

Author: vgvoleg

tests/query/test_query_transaction.py

@@ -1,5 +1,6 @@
 import pytest
 
+from ydb.query.transaction import QueryTxStateEnum
 
 class TestQueryTransaction:
     def test_tx_begin(self, tx):
@@ -46,3 +47,22 @@ def test_tx_execute_raises_after_rollback(self, tx):
         tx.rollback()
         with pytest.raises(RuntimeError):
             tx.execute("select 1;")
+
+    def test_context_manager_rollbacks_tx(self, tx):
+        with tx:
+            tx.begin()
+
+        assert tx._tx_state._state == QueryTxStateEnum.ROLLBACKED
+
+    def test_context_manager_normal_flow(self, tx):
+        with tx:
+            tx.begin()
+            tx.execute("select 1;")
+            tx.commit()
+
+        assert tx._tx_state._state == QueryTxStateEnum.COMMITTED
+
+    def test_context_manager_does_not_hide_exceptions(self, tx):
+        with pytest.raises(RuntimeError):
+            with tx:
+                tx.commit()

ydb/query/transaction.py

@@ -184,6 +184,7 @@ def __exit__(self, *args, **kwargs):
         Closes a transaction context manager and rollbacks transaction if
         it is not rolled back explicitly
         """
+        self._ensure_prev_stream_finished()
         if self._tx_state.tx_id is not None:
             # It's strictly recommended to close transactions directly
             # by using commit_tx=True flag while executing statement or by