Fix problem with specifying model files as argument (#514)

lib/annotate/annotate_models.rb

@@ -649,34 +649,52 @@ module AnnotateModels
     # of model files from root dir. Otherwise we take all the model files
     # in the model_dir directory.
     def get_model_files(options)
-      models = []
-      unless options[:is_rake]
-        models = ARGV.dup.reject { |m| m.match(/^(.*)=/) }
-      end
+      model_files = []
+
+      model_files = list_model_files_from_argument unless options[:is_rake]
+
+      return model_files unless model_files.empty?
 
-      if models.empty?
-        begin
       model_dir.each do |dir|
         Dir.chdir(dir) do
-              lst =
-                if options[:ignore_model_sub_dir]
-                  Dir["*.rb"].map{ |f| [dir, f] }
+          list = if options[:ignore_model_sub_dir]
+                   Dir["*.rb"].map { |f| [dir, f] }
                  else
-                  Dir["**/*.rb"].reject{ |f| f["concerns/"] }.map{ |f| [dir, f] }
+                   Dir["**/*.rb"].reject { |f| f["concerns/"] }.map { |f| [dir, f] }
                  end
-              models.concat(lst)
+          model_files.concat(list)
         end
       end
+
+      model_files
     rescue SystemCallError
       puts "No models found in directory '#{model_dir.join("', '")}'."
       puts "Either specify models on the command line, or use the --model-dir option."
       puts "Call 'annotate --help' for more info."
       exit 1
     end
+
+    def list_model_files_from_argument
+      return [] if ARGV.empty?
+
+      specified_files = ARGV.map { |file| File.expand_path(file) }
+
+      model_files = model_dir.flat_map do |dir|
+        absolute_dir_path = File.expand_path(dir)
+        specified_files
+          .find_all { |file| file.start_with?(absolute_dir_path) }
+          .map { |file| [dir, file.sub("#{absolute_dir_path}/", '')] }
+      end
+
+      if model_files.size != specified_files.size
+        puts "The specified file could not be found in directory '#{model_dir.join("', '")}'."
+        puts "Call 'annotate --help' for more info."
+        exit 1
       end
 
-      models
+      model_files
     end
+    private :list_model_files_from_argument
 
     # Retrieve the classes belonging to the model names we're asked to process
     # Check for namespaced models in subdirectories as well as models

spec/annotate/annotate_models_spec.rb

@@ -3,6 +3,7 @@ require File.dirname(__FILE__) + '/../spec_helper.rb'
 require 'annotate/annotate_models'
 require 'annotate/active_record_patch'
 require 'active_support/core_ext/string'
+require 'files'
 
 describe AnnotateModels do
   def mock_index(name, params = {})
@@ -951,6 +952,121 @@ EOS
     end
   end
 
+  describe '#get_model_files' do
+    subject { described_class.get_model_files(options) }
+
+    before do
+      ARGV.clear
+
+      described_class.model_dir = [model_dir]
+    end
+
+    context 'when `model_dir` is valid' do
+      let(:model_dir) do
+        Files do
+          file 'foo.rb'
+          dir 'bar' do
+            file 'baz.rb'
+            dir 'qux' do
+              file 'quux.rb'
+            end
+          end
+          dir 'concerns' do
+            file 'corge.rb'
+          end
+        end
+      end
+
+      context 'when the model files are not specified' do
+        context 'when no option is specified' do
+          let(:options) { {} }
+
+          it 'returns all model files under `model_dir` directory' do
+            is_expected.to contain_exactly(
+              [model_dir, 'foo.rb'],
+              [model_dir, File.join('bar', 'baz.rb')],
+              [model_dir, File.join('bar', 'qux', 'quux.rb')]
+            )
+          end
+        end
+
+        context 'when `ignore_model_sub_dir` option is enabled' do
+          let(:options) { { ignore_model_sub_dir: true } }
+
+          it 'returns model files just below `model_dir` directory' do
+            is_expected.to contain_exactly([model_dir, 'foo.rb'])
+          end
+        end
+      end
+
+      context 'when the model files are specified' do
+        let(:additional_model_dir) { 'additional_model' }
+        let(:model_files) do
+          [
+            File.join(model_dir, 'foo.rb'),
+            "./#{File.join(additional_model_dir, 'corge/grault.rb')}" # Specification by relative path
+          ]
+        end
+
+        before { ARGV.concat(model_files) }
+
+        context 'when no option is specified' do
+          let(:options) { {} }
+
+          context 'when all the specified files are in `model_dir` directory' do
+            before do
+              described_class.model_dir << additional_model_dir
+            end
+
+            it 'returns specified files' do
+              is_expected.to contain_exactly(
+                [model_dir, 'foo.rb'],
+                [additional_model_dir, 'corge/grault.rb']
+              )
+            end
+          end
+
+          context 'when a model file outside `model_dir` directory is specified' do
+            it 'exits with the status code' do
+              begin
+                subject
+                raise
+              rescue SystemExit => e
+                expect(e.status).to eq(1)
+              end
+            end
+          end
+        end
+
+        context 'when `is_rake` option is enabled' do
+          let(:options) { { is_rake: true } }
+
+          it 'returns all model files under `model_dir` directory' do
+            is_expected.to contain_exactly(
+              [model_dir, 'foo.rb'],
+              [model_dir, File.join('bar', 'baz.rb')],
+              [model_dir, File.join('bar', 'qux', 'quux.rb')]
+            )
+          end
+        end
+      end
+    end
+
+    context 'when `model_dir` is invalid' do
+      let(:model_dir) { '/not_exist_path' }
+      let(:options) { {} }
+
+      it 'exits with the status code' do
+        begin
+          subject
+          raise
+        rescue SystemExit => e
+          expect(e.status).to eq(1)
+        end
+      end
+    end
+  end
+
   describe '#get_model_class' do
     require 'tmpdir'