lib/rspec/hive/with_hive_connection.rb in rspec-hive-0.1.0 vs lib/rspec/hive/with_hive_connection.rb in rspec-hive-0.2.0

@@ -1,27 +1,33 @@
+require_relative 'exponential_backoff'
+
 module RSpec
   module Hive
     module WithHiveConnection
-      def hive
-        Hive.connector
-      end
-
-      def connection
-        @connection ||= hive.start_connection
-      end
-
       def self.included(mod)
         mod.before(:all) do
-          connection
+          ExponentialBackoff.retryable(on: Thrift::TransportException) do
+            connection
+          end
         end
 
         mod.before(:each) do
           connection.switch_database(DbName.random_name)
         end
 
         mod.after(:all) do
-          hive.stop_connection(connection) unless hive && connection
+          hive_connector.stop_connection(connection) if hive_connector && @connection
         end
+      end
+
+      def connection
+        @connection ||= hive_connector.start_connection
+      end
+
+      private
+
+      def hive_connector
+        ::RSpec::Hive.connector
       end
     end
   end
 end