diff --git a/test/requests/code_actions_formatting_test.rb b/test/requests/code_actions_formatting_test.rb
index 8bcb493b01..7d85cfd4d6 100644
--- a/test/requests/code_actions_formatting_test.rb
+++ b/test/requests/code_actions_formatting_test.rb
@@ -83,7 +83,8 @@ def assert_corrects_to_expected(diagnostic_code, code_action_title, source, expe
     )
 
     diagnostics = RubyLsp::Requests::Diagnostics.new(global_state, document).perform
-    diagnostic = T.must(T.must(diagnostics).find { |d| d.code == diagnostic_code })
+    # binding.irb
+    diagnostic = T.must(T.must(diagnostics).find { |d| d.attributes[:code] && (d.code == diagnostic_code) })
     range = diagnostic.range.to_hash.transform_values(&:to_hash)
     result = RubyLsp::Requests::CodeActions.new(document, range, {
       diagnostics: [JSON.parse(T.must(diagnostic).to_json, symbolize_names: true)],