From da5a2521fc7a063f376c20d015ba94e948c1de67 Mon Sep 17 00:00:00 2001
From: Oleg Pudeyev <oleg@bsdpower.com>
Date: Mon, 1 Feb 2021 21:09:10 -0500
Subject: [PATCH] RUBY-2484 make it possible to request that sessions are not
 killed to work around atlas not allowing their death

---
 spec/runners/unified/test.rb | 34 +++++++++++++++++++++-------------
 1 file changed, 21 insertions(+), 13 deletions(-)

diff --git a/spec/runners/unified/test.rb b/spec/runners/unified/test.rb
index cb96e120c4..228b923ebc 100644
--- a/spec/runners/unified/test.rb
+++ b/spec/runners/unified/test.rb
@@ -187,6 +187,8 @@ def set_initial_data
     end
 
     def run
+      kill_sessions
+
       test_spec = UsingHash[self.test_spec]
       ops = test_spec.use!('operations')
       execute_operations(ops)
@@ -207,19 +209,7 @@ def stop?
 
     def cleanup
       if $kill_transactions || true
-        begin
-          root_authorized_client.command(
-            killAllSessions: [],
-          )
-        rescue Mongo::Error::OperationFailure => e
-          if e.code == 11601
-            # operation was interrupted, ignore
-          elsif e.code == 59
-            # no such command (old server), ignore
-          else
-            raise
-          end
-        end
+        kill_sessions
         $kill_transactions = nil
       end
 
@@ -333,6 +323,24 @@ def disable_fail_points
       end
     end
 
+    def kill_sessions
+      if options[:kill_sessions] != false
+        begin
+          root_authorized_client.command(
+            killAllSessions: [],
+          )
+        rescue Mongo::Error::OperationFailure => e
+          if e.code == 11601
+            # operation was interrupted, ignore
+          elsif e.code == 59
+            # no such command (old server), ignore
+          else
+            raise
+          end
+        end
+      end
+    end
+
     def root_authorized_client
       @root_authorized_client ||= ClientRegistry.instance.global_client('root_authorized')
     end