1 files changed, 42 insertions, 28 deletions
diff --git a/railties/test/application/rake/notes_test.rb b/railties/test/application/rake/notes_test.rb
index 50def9beb0..8e9fe9b6b4 100644
--- a/railties/test/application/rake/notes_test.rb
+++ b/railties/test/application/rake/notes_test.rb
@@ -1,5 +1,7 @@
+# frozen_string_literal: true
+
 require "isolation/abstract_unit"
-require 'rails/source_annotation_extractor'
+require "rails/source_annotation_extractor"
 
 module ApplicationTests
   module RakeTests
@@ -17,15 +19,15 @@ module ApplicationTests
         teardown_app
       end
 
-      test 'notes finds notes for certain file_types' do
+      test "notes finds notes for certain file_types" do
         app_file "app/views/home/index.html.erb", "<% # TODO: note in erb %>"
         app_file "app/assets/javascripts/application.js", "// TODO: note in js"
         app_file "app/assets/stylesheets/application.css", "// TODO: note in css"
         app_file "app/controllers/application_controller.rb", 1000.times.map { "" }.join("\n") << "# TODO: note in ruby"
         app_file "lib/tasks/task.rake", "# TODO: note in rake"
-        app_file 'app/views/home/index.html.builder', '# TODO: note in builder'
-        app_file 'config/locales/en.yml', '# TODO: note in yml'
-        app_file 'config/locales/en.yaml', '# TODO: note in yaml'
+        app_file "app/views/home/index.html.builder", "# TODO: note in builder"
+        app_file "config/locales/en.yml", "# TODO: note in yml"
+        app_file "config/locales/en.yaml", "# TODO: note in yaml"
         app_file "app/views/home/index.ruby", "# TODO: note in ruby"
 
         run_rake_notes do |output, lines|
@@ -43,7 +45,7 @@ module ApplicationTests
         end
       end
 
-      test 'notes finds notes in default directories' do
+      test "notes finds notes in default directories" do
         app_file "app/controllers/some_controller.rb", "# TODO: note in app directory"
         app_file "config/initializers/some_initializer.rb", "# TODO: note in config directory"
         app_file "db/some_seeds.rb", "# TODO: note in db directory"
@@ -65,7 +67,7 @@ module ApplicationTests
         end
       end
 
-      test 'notes finds notes in custom directories' do
+      test "notes finds notes in custom directories" do
         app_file "app/controllers/some_controller.rb", "# TODO: note in app directory"
         app_file "config/initializers/some_initializer.rb", "# TODO: note in config directory"
         app_file "db/some_seeds.rb", "# TODO: note in db directory"
@@ -88,9 +90,10 @@ module ApplicationTests
         end
       end
 
-      test 'custom rake task finds specific notes in specific directories' do
+      test "custom rake task finds specific notes in specific directories" do
         app_file "app/controllers/some_controller.rb", "# TODO: note in app directory"
-        app_file "lib/some_file.rb", "# OPTIMIZE: note in lib directory\n" << "# FIXME: note in lib directory"
+        app_file "lib/some_file.rb", "# OPTIMIZE: note in lib directory\n" \
+          "# FIXME: note in lib directory"
         app_file "test/some_test.rb", 1000.times.map { "" }.join("\n") << "# TODO: note in test directory"
 
         app_file "lib/tasks/notes_custom.rake", <<-EOS
@@ -113,7 +116,7 @@ module ApplicationTests
         end
       end
 
-      test 'register a new extension' do
+      test "register a new extension" do
         add_to_config "config.assets.precompile = []"
         add_to_config %q{ config.annotations.register_extensions("scss", "sass") { |annotation| /\/\/\s*(#{annotation}):?\s*(.*)$/ } }
         app_file "app/assets/stylesheets/application.css.scss", "// TODO: note in scss"
@@ -126,32 +129,43 @@ module ApplicationTests
         end
       end
 
+      test "register additional directories" do
+        app_file "spec/spec_helper.rb", "# TODO: note in spec"
+        app_file "spec/models/user_spec.rb", "# TODO: note in model spec"
+        add_to_config ' config.annotations.register_directories("spec") '
+
+        run_rake_notes do |output, lines|
+          assert_match(/note in spec/, output)
+          assert_match(/note in model spec/, output)
+          assert_equal 2, lines.size
+        end
+      end
+
       private
 
-      def run_rake_notes(command = 'bin/rails notes')
-        boot_rails
-        load_tasks
+        def run_rake_notes(command = "bin/rails notes")
+          boot_rails
+          load_tasks
 
-        Dir.chdir(app_path) do
-          output = `#{command}`
-          lines  = output.scan(/\[([0-9\s]+)\]\s/).flatten
+          Dir.chdir(app_path) do
+            output = `#{command}`
+            lines  = output.scan(/\[([0-9\s]+)\]\s/).flatten
 
-          yield output, lines
+            yield output, lines
+          end
         end
-      end
 
-      def load_tasks
-        require 'rake'
-        require 'rdoc/task'
-        require 'rake/testtask'
+        def load_tasks
+          require "rake"
+          require "rdoc/task"
+          require "rake/testtask"
 
-        Rails.application.load_tasks
-      end
+          Rails.application.load_tasks
+        end
 
-      def boot_rails
-        super
-        require "#{app_path}/config/environment"
-      end
+        def boot_rails
+          require "#{app_path}/config/environment"
+        end
     end
   end
 end