From ae4bd5cef32e74d2ff30bf146f8946d683369d88 Mon Sep 17 00:00:00 2001
From: Alex Averbuch <alex.averbuch@gmail.com>
Date: Wed, 23 Aug 2017 14:48:20 +0200
Subject: [PATCH] Propagates correct error status from test.loop procedure
 Reduces chance of race in flaky test

---
 .../auth/BuiltInProceduresInteractionTestBase.java   |  7 ++++++-
 .../auth/ProcedureInteractionTestBase.java           | 12 +++++++++---
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/enterprise/security/src/test/java/org/neo4j/server/security/enterprise/auth/BuiltInProceduresInteractionTestBase.java b/enterprise/security/src/test/java/org/neo4j/server/security/enterprise/auth/BuiltInProceduresInteractionTestBase.java
index c9521cf93152e..b4b4aa305106c 100644
--- a/enterprise/security/src/test/java/org/neo4j/server/security/enterprise/auth/BuiltInProceduresInteractionTestBase.java
+++ b/enterprise/security/src/test/java/org/neo4j/server/security/enterprise/auth/BuiltInProceduresInteractionTestBase.java
@@ -699,7 +699,9 @@ public void shouldTerminateLongRunningProcedureThatChecksTheGuardRegularlyIfKill
 
         String loopQuery = "CALL test.loop";
 
-        new Thread( () -> assertFail( readSubject, loopQuery, "Explicitly terminated by the user." ) ).start();
+        Thread loopQueryThread =
+                new Thread( () -> assertFail( readSubject, loopQuery, "Explicitly terminated by the user." ) );
+        loopQueryThread.start();
         latch.startAndWaitForAllToStart();
 
         try
@@ -725,6 +727,9 @@ public void shouldTerminateLongRunningProcedureThatChecksTheGuardRegularlyIfKill
             latch.finishAndWaitForAllToFinish();
         }
 
+        // there is a race with "test.loop" procedure - after decrementing latch it may take time to actually exit
+        loopQueryThread.join( 10_000 );
+
         assertEmpty(
                 adminSubject,
                 "CALL dbms.listQueries() YIELD query WITH * WHERE NOT query CONTAINS 'listQueries' RETURN *" );
diff --git a/enterprise/security/src/test/java/org/neo4j/server/security/enterprise/auth/ProcedureInteractionTestBase.java b/enterprise/security/src/test/java/org/neo4j/server/security/enterprise/auth/ProcedureInteractionTestBase.java
index 42e2e87058528..eddfe8f919dc0 100644
--- a/enterprise/security/src/test/java/org/neo4j/server/security/enterprise/auth/ProcedureInteractionTestBase.java
+++ b/enterprise/security/src/test/java/org/neo4j/server/security/enterprise/auth/ProcedureInteractionTestBase.java
@@ -59,7 +59,6 @@
 import org.neo4j.kernel.api.bolt.BoltConnectionTracker;
 import org.neo4j.kernel.api.bolt.ManagedBoltStateMachine;
 import org.neo4j.kernel.api.exceptions.InvalidArgumentsException;
-import org.neo4j.kernel.api.exceptions.Status;
 import org.neo4j.kernel.enterprise.builtinprocs.EnterpriseBuiltInDbmsProcedures;
 import org.neo4j.kernel.impl.proc.Procedures;
 import org.neo4j.logging.Log;
@@ -679,9 +678,16 @@ public void loop()
                     guard.check();
                 }
             }
-            catch (TransactionTerminatedException | TransactionGuardException e)
+            catch ( TransactionTerminatedException | TransactionGuardException e )
             {
-                throw new TransactionGuardException( TransactionTimedOut, PROCEDURE_TIMEOUT_ERROR, e );
+                if ( e.status().equals( TransactionTimedOut ) )
+                {
+                    throw new TransactionGuardException( TransactionTimedOut, PROCEDURE_TIMEOUT_ERROR, e );
+                }
+                else
+                {
+                    throw e;
+                }
             }
             finally
             {