1 files changed, 75 insertions, 0 deletions
diff --git a/spec/rubocop/code_reuse_helpers_spec.rb b/spec/rubocop/code_reuse_helpers_spec.rb
index 695c152e3db..3220cff1681 100644
--- a/spec/rubocop/code_reuse_helpers_spec.rb
+++ b/spec/rubocop/code_reuse_helpers_spec.rb
@@ -21,6 +21,8 @@ RSpec.describe RuboCop::CodeReuseHelpers do
     end.new
   end
 
+  let(:ee_file_path) { File.expand_path('../../ee/app/models/license.rb', __dir__) }
+
   describe '#send_to_constant?' do
     it 'returns true when sending to a constant' do
       node = build_and_parse_source('Foo.bar')
@@ -312,4 +314,77 @@ RSpec.describe RuboCop::CodeReuseHelpers do
       cop.disallow_send_to(def_node, 'Finder', 'oops')
     end
   end
+
+  describe '#ee?' do
+    before do
+      stub_env('FOSS_ONLY', nil)
+      allow(File).to receive(:exist?).with(ee_file_path) { true }
+    end
+
+    it 'returns true when ee/app/models/license.rb exists' do
+      expect(cop.ee?).to eq(true)
+    end
+  end
+
+  describe '#jh?' do
+    context 'when jh directory exists and EE_ONLY is not set' do
+      before do
+        stub_env('EE_ONLY', nil)
+
+        allow(Dir).to receive(:exist?).with(File.expand_path('../../jh', __dir__)) { true }
+      end
+
+      context 'when ee/app/models/license.rb exists' do
+        before do
+          allow(File).to receive(:exist?).with(ee_file_path) { true }
+        end
+
+        context 'when FOSS_ONLY is not set' do
+          before do
+            stub_env('FOSS_ONLY', nil)
+          end
+
+          it 'returns true' do
+            expect(cop.jh?).to eq(true)
+          end
+        end
+
+        context 'when FOSS_ONLY is set to 1' do
+          before do
+            stub_env('FOSS_ONLY', '1')
+          end
+
+          it 'returns false' do
+            expect(cop.jh?).to eq(false)
+          end
+        end
+      end
+
+      context 'when ee/app/models/license.rb not exist' do
+        before do
+          allow(File).to receive(:exist?).with(ee_file_path) { false }
+        end
+
+        context 'when FOSS_ONLY is not set' do
+          before do
+            stub_env('FOSS_ONLY', nil)
+          end
+
+          it 'returns true' do
+            expect(cop.jh?).to eq(false)
+          end
+        end
+
+        context 'when FOSS_ONLY is set to 1' do
+          before do
+            stub_env('FOSS_ONLY', '1')
+          end
+
+          it 'returns false' do
+            expect(cop.jh?).to eq(false)
+          end
+        end
+      end
+    end
+  end
 end