When copying ServiceTemplate copy only direct_custom_buttons

app/models/service_template/copy.rb

@@ -7,7 +7,7 @@ def template_copy(new_name = "Copy of " + name + Time.zone.now.to_s)
         dup.tap do |template|
           template.update_attributes(:name => new_name, :display => false)
           service_resources.each { |service_resource| resource_copy(service_resource, template) }
-          custom_buttons.each { |custom_button| custom_button_copy(custom_button, template) }
+          direct_custom_buttons.each { |custom_button| custom_button_copy(custom_button, template) }
           custom_button_sets.each { |custom_button_set| custom_button_set_copy(custom_button_set, template) }
         end.save!
       end

spec/models/service_template_spec.rb

@@ -272,7 +272,9 @@
   context "#template_copy" do
     let(:service_template_ansible_tower) { FactoryBot.create(:service_template_ansible_tower) }
     let(:service_template_orchestration) { FactoryBot.create(:service_template_orchestration) }
-    let(:custom_button) { FactoryBot.create(:custom_button, :applies_to_class => "Service") }
+    let(:service) {FactoryBot.create(:service, :service_template_id => @st1.id)}
+    let(:custom_button) { FactoryBot.create(:custom_button, :applies_to => @st1) }
+    let(:custom_button_for_service) { FactoryBot.create(:custom_button, :applies_to_class => "Service") }
     let(:custom_button_set) { FactoryBot.create(:custom_button_set, :owner => @st1) }
     before do
       @st1 = FactoryBot.create(:service_template)
@@ -288,15 +290,17 @@
         expect(new_service_template.guid).not_to eq(@st1.guid)
       end
 
-      it "with custom button" do
+      it "with custom button copy only direct_custom_buttons" do
         custom_button
-        expect(@st1.custom_buttons.count).to eq(1)
+        custom_button_for_service
+        expect(@st1.custom_buttons.count).to eq(2)
+        number_of_service_templates = ServiceTemplate.count
         @st1.template_copy("new_template")
-        expect(ServiceTemplate.count).to eq(2)
+        expect(ServiceTemplate.count).to eq(number_of_service_templates + 1)
         new_service_template = ServiceTemplate.find_by(:name => "new_template")
         expect(new_service_template.display).to be(false)
         expect(new_service_template.guid).not_to eq(@st1.guid)
-        expect(new_service_template.custom_buttons.count).to eq(2)
+        expect(new_service_template.direct_custom_buttons.count).to eq(@st1.direct_custom_buttons.count)
       end
 
       it "with custom button set" do