diff --git a/features/steps/shared/diff_note.rb b/features/steps/shared/diff_note.rb
index f80d8d064755d50228d44f0792e9acff139c717b..f917d7bde08c1f284200cca60f4ec5c88a6a5988 100644
--- a/features/steps/shared/diff_note.rb
+++ b/features/steps/shared/diff_note.rb
@@ -138,7 +138,7 @@ module SharedDiffNote
 
   Then 'I should see the diff comment edit button' do
     within(".diff-file") do
-      page.should have_css(".js-note-edit-button", visible: true)
+      page.should have_css(".js-note-write-button", visible: true)
     end
   end
 
diff --git a/features/steps/shared/note.rb b/features/steps/shared/note.rb
index 36b81b74186dc7e0a1ebd83f15ca5833304d2817..1c52d4c72d89a7ce4da56a4a0186cc194dfb4002 100644
--- a/features/steps/shared/note.rb
+++ b/features/steps/shared/note.rb
@@ -81,7 +81,7 @@ module SharedNote
 
   Then 'I should see the comment edit button' do
     within(".js-main-target-form") do
-      page.should have_css(".js-note-edit-button", visible: true)
+      page.should have_css(".js-note-write-button", visible: true)
     end
   end