spec/mongo/retryable_spec.rb in mongo-2.17.4 vs spec/mongo/retryable_spec.rb in mongo-2.18.0.beta1

@@ -37,11 +37,12 @@
   end
 
   def write
     # This passes a nil session and therefore triggers
     # legacy_write_with_retry code path
-    write_with_retry(session, write_concern) do
+    context = Mongo::Operation::Context.new(client: @client, session: session)
+    write_with_retry(write_concern, context: context) do
       operation.execute
     end
   end
 
   def retry_write_allowed_as_configured?
@@ -63,17 +64,20 @@
   include RSpec::Mocks::ExampleMethods
 
   def session
     double('session').tap do |session|
       expect(session).to receive(:retry_writes?).and_return(true)
+      allow(session).to receive(:materialize_if_needed)
 
       # mock everything else that is in the way
       i = 1
       allow(session).to receive(:next_txn_num) { i += 1 }
       allow(session).to receive(:in_transaction?).and_return(false)
       allow(session).to receive(:pinned_server)
+      allow(session).to receive(:pinned_connection_global_id)
       allow(session).to receive(:starting_transaction?).and_return(false)
+      allow(session).to receive(:materialize)
     end
   end
 
   def write_concern
     nil
@@ -90,12 +94,20 @@
 
   let(:operation) do
     double('operation')
   end
 
-  let(:server) { double('server') }
+  let(:connection) do
+    double('connection')
+  end
 
+  let(:server) do
+    double('server').tap do |server|
+      allow(server).to receive('with_connection').and_yield(connection)
+    end
+  end
+
   let(:max_read_retries) { 1 }
   let(:max_write_retries) { 1 }
 
   let(:cluster) do
     double('cluster', next_primary: server).tap do |cluster|
@@ -118,10 +130,21 @@
 
   let(:retryable) do
     LegacyRetryableTestConsumer.new(operation, cluster, client)
   end
 
+  let(:session) do
+    double('session').tap do |session|
+      allow(session).to receive(:pinned_connection_global_id).and_return(nil)
+      allow(session).to receive(:materialize_if_needed)
+    end
+  end
+
+  let(:context) do
+    Mongo::Operation::Context.new(client: client, session: session)
+  end
+
   before do
     # Retryable reads perform server selection
     allow_any_instance_of(Mongo::ServerSelector::Primary).to receive(:select_server).and_return(server)
   end
 
@@ -139,12 +162,18 @@
     end
 
     context 'when ending_transaction is true' do
       let(:retryable) { RetryableTestConsumer.new(operation, cluster, client) }
 
+      let(:context) do
+        Mongo::Operation::Context.new(client: client, session: nil)
+      end
+
       it 'raises ArgumentError' do
         expect do
-          retryable.write_with_retry(nil, nil, true)
+          retryable.write_with_retry(nil, ending_transaction: true, context: context) do
+            fail 'Expected not to get here'
+          end
         end.to raise_error(ArgumentError, 'Cannot end a transaction without a session')
       end
     end
 
     context 'when a socket error occurs' do