Tests for allowing more statements and tx after reset

Author: Zhen Li

driver/src/main/java/org/neo4j/driver/internal/NetworkSession.java

@@ -59,7 +59,6 @@ public void run()
 
     private ExplicitTransaction currentTransaction;
     private AtomicBoolean isOpen = new AtomicBoolean( true );
-    private AtomicBoolean markedToClose = new AtomicBoolean( false );
 
     NetworkSession( Connection connection, Logger logger )
     {
@@ -106,19 +105,17 @@ public StatementResult run( Statement statement )
 
     public void reset()
     {
+        ensureSessionIsOpen();
         ensureNoUnrecoverableError();
         ensureConnectionIsOpen();
 
-        if( markedToClose.compareAndSet( false, true ) )
-        {
-            connection.resetAsync();
-        }
+        connection.resetAsync();
     }
 
     @Override
     public boolean isOpen()
     {
-        return isOpen.get() && !markedToClose.get();
+        return isOpen.get();
     }
 
     @Override
@@ -206,13 +203,15 @@ public TypeSystem typeSystem()
 
     private void ensureConnectionIsValidBeforeRunningSession()
     {
+        ensureSessionIsOpen();
         ensureNoUnrecoverableError();
         ensureNoOpenTransactionBeforeRunningSession();
         ensureConnectionIsOpen();
     }
 
     private void ensureConnectionIsValidBeforeOpeningTransaction()
     {
+        ensureSessionIsOpen();
         ensureNoUnrecoverableError();
         ensureNoOpenTransactionBeforeOpeningTransaction();
         ensureConnectionIsOpen();

driver/src/test/java/org/neo4j/driver/v1/integration/SessionIT.java

@@ -21,12 +21,7 @@
 import org.junit.Rule;
 import org.junit.Test;
 
-import org.neo4j.driver.v1.AuthToken;
-import org.neo4j.driver.v1.AuthTokens;
-import org.neo4j.driver.v1.Driver;
-import org.neo4j.driver.v1.GraphDatabase;
-import org.neo4j.driver.v1.Session;
-import org.neo4j.driver.v1.StatementResult;
+import org.neo4j.driver.v1.*;
 import org.neo4j.driver.v1.exceptions.ClientException;
 import org.neo4j.driver.v1.exceptions.Neo4jException;
 import org.neo4j.driver.v1.util.TestNeo4j;
@@ -126,7 +121,6 @@ public void shouldKillLongRunningStatement() throws Throwable
             assertTrue( startTime > 0 );
             assertTrue( endTime - startTime > killTimeout * 1000 ); // get reset by session.reset
             assertTrue( endTime - startTime < executionTimeout * 1000 / 2 ); // finished before execution finished
-            assertFalse( session.isOpen() );
         }
         catch ( Exception e )
         {
@@ -201,4 +195,47 @@ public void run()
             }
         } ).start();
     }
+
+    @Test
+    public void shouldAllowMoreStatementAfterSessionReset()
+    {
+        // Given
+        try( Driver driver =  GraphDatabase.driver( neo4j.uri() );
+             Session session = driver.session() )
+        {
+
+            session.run( "Return 1" ).consume();
+
+            // When reset the state of this session
+            session.reset();
+
+            // Then can run successfully more statements without any error
+            session.run( "Return 2" ).consume();
+        }
+    }
+
+    @Test
+    public void shouldAllowMoreTxAfterSessionReset()
+    {
+        // Given
+        try( Driver driver =  GraphDatabase.driver( neo4j.uri() );
+             Session session = driver.session() )
+        {
+            try( Transaction tx = session.beginTransaction() )
+            {
+                tx.run("Return 1");
+                tx.success();
+            }
+
+            // When reset the state of this session
+            session.reset();
+
+            // Then can run more Tx
+            try( Transaction tx = session.beginTransaction() )
+            {
+                tx.run("Return 2");
+                tx.success();
+            }
+        }
+    }
 }