1 files changed, 73 insertions, 8 deletions

diff --git a/spec/support/shared_examples/namespaces/traversal_scope_examples.rb b/spec/support/shared_examples/namespaces/traversal_scope_examples.rb
index 74b1bacc560..4c09c1c2a3b 100644
--- a/spec/support/shared_examples/namespaces/traversal_scope_examples.rb
+++ b/spec/support/shared_examples/namespaces/traversal_scope_examples.rb
@@ -25,12 +25,6 @@ RSpec.shared_examples 'namespace traversal scopes' do
     it { is_expected.to contain_exactly(group_1.id, group_2.id) }
   end
 
-  describe '.without_sti_condition' do
-    subject { described_class.without_sti_condition }
-
-    it { expect(subject.where_values_hash).not_to have_key(:type) }
-  end
-
   describe '.order_by_depth' do
     subject { described_class.where(id: [group_1, nested_group_1, deep_nested_group_1]).order_by_depth(direction) }
 
@@ -55,6 +49,53 @@ RSpec.shared_examples 'namespace traversal scopes' do
     it { is_expected.to eq described_class.column_names }
   end
 
+  shared_examples '.roots' do
+    context 'with only sub-groups' do
+      subject { described_class.where(id: [deep_nested_group_1, nested_group_1, deep_nested_group_2]).roots }
+
+      it { is_expected.to contain_exactly(group_1, group_2) }
+    end
+
+    context 'with only root groups' do
+      subject { described_class.where(id: [group_1, group_2]).roots }
+
+      it { is_expected.to contain_exactly(group_1, group_2) }
+    end
+
+    context 'with all groups' do
+      subject { described_class.where(id: groups).roots }
+
+      it { is_expected.to contain_exactly(group_1, group_2) }
+    end
+  end
+
+  describe '.roots' do
+    context "use_traversal_ids_roots feature flag is true" do
+      before do
+        stub_feature_flags(use_traversal_ids: true)
+        stub_feature_flags(use_traversal_ids_roots: true)
+      end
+
+      it_behaves_like '.roots'
+
+      it 'not make recursive queries' do
+        expect { described_class.where(id: [nested_group_1]).roots.load }.not_to make_queries_matching(/WITH RECURSIVE/)
+      end
+    end
+
+    context "use_traversal_ids_roots feature flag is false" do
+      before do
+        stub_feature_flags(use_traversal_ids_roots: false)
+      end
+
+      it_behaves_like '.roots'
+
+      it 'make recursive queries' do
+        expect { described_class.where(id: [nested_group_1]).roots.load }.to make_queries_matching(/WITH RECURSIVE/)
+      end
+    end
+  end
+
   shared_examples '.self_and_ancestors' do
     subject { described_class.where(id: [nested_group_1, nested_group_2]).self_and_ancestors }
 
@@ -156,7 +197,7 @@ RSpec.shared_examples 'namespace traversal scopes' do
     end
   end
 
-  describe '.self_and_descendants' do
+  shared_examples '.self_and_descendants' do
     subject { described_class.where(id: [nested_group_1, nested_group_2]).self_and_descendants }
 
     it { is_expected.to contain_exactly(nested_group_1, deep_nested_group_1, nested_group_2, deep_nested_group_2) }
@@ -174,7 +215,19 @@ RSpec.shared_examples 'namespace traversal scopes' do
     end
   end
 
-  describe '.self_and_descendant_ids' do
+  describe '.self_and_descendants' do
+    include_examples '.self_and_descendants'
+
+    context 'with traversal_ids_btree feature flag disabled' do
+      before do
+        stub_feature_flags(traversal_ids_btree: false)
+      end
+
+      include_examples '.self_and_descendants'
+    end
+  end
+
+  shared_examples '.self_and_descendant_ids' do
     subject { described_class.where(id: [nested_group_1, nested_group_2]).self_and_descendant_ids.pluck(:id) }
 
     it { is_expected.to contain_exactly(nested_group_1.id, deep_nested_group_1.id, nested_group_2.id, deep_nested_group_2.id) }
@@ -190,4 +243,16 @@ RSpec.shared_examples 'namespace traversal scopes' do
       it { is_expected.to contain_exactly(deep_nested_group_1.id, deep_nested_group_2.id) }
     end
   end
+
+  describe '.self_and_descendant_ids' do
+    include_examples '.self_and_descendant_ids'
+
+    context 'with traversal_ids_btree feature flag disabled' do
+      before do
+        stub_feature_flags(traversal_ids_btree: false)
+      end
+
+      include_examples '.self_and_descendant_ids'
+    end
+  end
 end