Include test files in test file result

lib/quality/test_file_finder.rb

 # frozen_string_literal: true
 
 require 'ostruct'
+require 'set'
 
 module Quality
   class TestFileFinder
@@ -10,6 +11,7 @@ module Quality
     def initialize(file, foss_test_only: false)
       @file = file
       @foss_test_only = foss_test_only
+      @result = Set.new
     end
 
     def test_files
@@ -17,18 +19,22 @@ module Quality
       contexts.flat_map do |context|
         match_test_files_for(context)
       end
+
+      result.to_a
     end
 
     private
 
-    attr_reader :file, :foss_test_only
+    attr_reader :file, :foss_test_only, :result
 
     def ee_context
       OpenStruct.new.tap do |ee|
         ee.app = %r{^#{EE_PREFIX}app/(.+)\.rb$} unless foss_test_only
         ee.lib = %r{^#{EE_PREFIX}lib/(.+)\.rb$} unless foss_test_only
+        ee.spec = %r{^#{EE_PREFIX}spec/(.+)_spec.rb$} unless foss_test_only
         ee.spec_dir = "#{EE_PREFIX}spec" unless foss_test_only
-        ee.ee_modules = %r{^#{EE_PREFIX}(.*\/)ee/(.+)\.rb$}
+        ee.ee_modules = %r{^#{EE_PREFIX}(?!spec)(.*\/)ee/(.+)\.rb$}
+        ee.ee_module_spec = %r{^#{EE_PREFIX}spec/(.*\/)ee/(.+)\.rb$}
         ee.foss_spec_dir = 'spec'
       end
     end
@@ -37,26 +43,31 @@ module Quality
       OpenStruct.new.tap do |foss|
         foss.app = %r{^app/(.+)\.rb$}
         foss.lib = %r{^lib/(.+)\.rb$}
+        foss.spec = %r{^spec/(.+)_spec.rb$}
         foss.spec_dir = 'spec'
       end
     end
 
     def match_test_files_for(context)
-      test_files = []
-
       if (match = context.app&.match(file))
-        test_files << "#{context.spec_dir}/#{match[1]}_spec.rb"
+        result << "#{context.spec_dir}/#{match[1]}_spec.rb"
       end
 
       if (match = context.lib&.match(file))
-        test_files << "#{context.spec_dir}/lib/#{match[1]}_spec.rb"
+        result << "#{context.spec_dir}/lib/#{match[1]}_spec.rb"
+      end
+
+      if context.spec&.match(file)
+        result << file
       end
 
       if (match = context.ee_modules&.match(file))
-        test_files << "#{context.foss_spec_dir}/#{match[1]}#{match[2]}_spec.rb"
+        result << "#{context.foss_spec_dir}/#{match[1]}#{match[2]}_spec.rb"
       end
 
-      test_files
+      if (match = context.ee_module_spec&.match(file))
+        result << "#{context.foss_spec_dir}/#{match[1]}#{match[2]}.rb"
+      end
     end
   end
 end

spec/lib/quality/test_file_finder_spec.rb

@@ -45,8 +45,9 @@ describe Quality::TestFileFinder do
 
     context 'when given a test file' do
       let(:file) { 'spec/lib/banzai/color_parser_spec.rb' }
+      let(:test_files) { ['spec/lib/banzai/color_parser_spec.rb'] }
 
-      it_behaves_like 'not finding a matching test file'
+      it_behaves_like 'finding matching test files'
     end
 
     context 'when given an ee app file' do
@@ -71,9 +72,17 @@ describe Quality::TestFileFinder do
     end
 
     context 'when given an ee test file' do
-      let(:file) { 'ee/spec/lib/banzai/color_parser_spec.rb' }
+      let(:file) { 'ee/spec/models/container_registry/event_spec.rb' }
+      let(:test_files) { ['ee/spec/models/container_registry/event_spec.rb'] }
 
-      it_behaves_like 'not finding a matching test file'
+      it_behaves_like 'finding matching test files'
+    end
+
+    context 'when given an ee module test file' do
+      let(:file) { 'ee/spec/models/ee/appearance_spec.rb' }
+      let(:test_files) { ['ee/spec/models/ee/appearance_spec.rb', 'spec/models/appearance_spec.rb'] }
+
+      it_behaves_like 'finding matching test files'
     end
 
     context 'with foss_test_only: true' do