Ensure modified documents are cleared

lib/mongoid/clients/sessions.rb

@@ -62,6 +62,7 @@ def with_session(options = {})
         rescue *transactions_not_supported_exceptions
           raise Mongoid::Errors::TransactionsNotSupported
         ensure
+          Threaded.clear_modified_documents(session)
           Threaded.clear_session(client: persistence_context.client)
         end
 

lib/mongoid/threaded.rb

@@ -18,6 +18,7 @@ module Threaded
     CURRENT_SCOPE_KEY = '[mongoid]:current-scope'
 
     AUTOSAVES_KEY = '[mongoid]:autosaves'
+
     VALIDATIONS_KEY = '[mongoid]:validations'
 
     STACK_KEYS = Hash.new do |hash, key|
@@ -376,9 +377,7 @@ def add_modified_document(session, document)
     # @return [ Set<Mongoid::Document> ] Collection of modified documents before
     #   it was cleared.
     def clear_modified_documents(session)
-      modified_documents[session].dup
-    ensure
-      modified_documents.delete(session)
+      modified_documents.delete(session) || []
     end
 
     # Queries whether document callbacks should be executed by default for the

spec/mongoid/threaded_spec.rb

@@ -343,25 +343,20 @@
   end
 
   describe '#clear_modified_documents' do
-    around do |example|
-      session = Mongoid.default_client.start_session
-      session.start_transaction
-      doc = Minim.create!
-      described_class.add_modified_document(session, doc)
-      described_class.clear_modified_documents(session)
-      example.run(session)
-    ensure
-      session.abort_transaction
-      session.end_session
+    let(:session) do
+      double(Mongo::Session).tap do |session|
+        allow(session).to receive(:in_transaction?).and_return(true)
+      end
     end
 
     context 'when there are modified documents' do
-      it 'removes the documents' do
-        expect(described_class.modified_documents).to be_empty
+      before do
+        described_class.add_modified_document(session, Minim.new)
+        described_class.clear_modified_documents(session)
       end
 
-      it 'removes the key' do |session|
-        expect(described_class.modified_documents.key?(session)).to be_falsey
+      it 'removes the documents and keys' do
+        expect(described_class.modified_documents).to be_empty
       end
     end
   end