2 files changed, 39 insertions, 1 deletions

diff --git a/spec/frontend/ci/pipeline_editor/components/editor/ci_editor_header_spec.js b/spec/frontend/ci/pipeline_editor/components/editor/ci_editor_header_spec.js
index 2861fc35342..f1a5c4169fb 100644
--- a/spec/frontend/ci/pipeline_editor/components/editor/ci_editor_header_spec.js
+++ b/spec/frontend/ci/pipeline_editor/components/editor/ci_editor_header_spec.js
@@ -11,12 +11,25 @@ describe('CI Editor Header', () => {
   let wrapper;
   let trackingSpy = null;
 
-  const createComponent = ({ showDrawer = false, showJobAssistantDrawer = false } = {}) => {
+  const createComponent = ({
+    showDrawer = false,
+    showJobAssistantDrawer = false,
+    showAiAssistantDrawer = false,
+    aiChatAvailable = false,
+    aiCiConfigGenerator = false,
+  } = {}) => {
     wrapper = extendedWrapper(
       shallowMount(CiEditorHeader, {
+        provide: {
+          aiChatAvailable,
+          glFeatures: {
+            aiCiConfigGenerator,
+          },
+        },
         propsData: {
           showDrawer,
           showJobAssistantDrawer,
+          showAiAssistantDrawer,
         },
       }),
     );
@@ -24,6 +37,7 @@ describe('CI Editor Header', () => {
 
   const findLinkBtn = () => wrapper.findByTestId('template-repo-link');
   const findHelpBtn = () => wrapper.findByTestId('drawer-toggle');
+  const findAiAssistnantBtn = () => wrapper.findByTestId('ai-assistant-drawer-toggle');
 
   afterEach(() => {
     unmockTracking();
@@ -39,7 +53,29 @@ describe('CI Editor Header', () => {
       label,
     });
   };
+  describe('Ai Assistant toggle button', () => {
+    describe('when feature is unavailable', () => {
+      it('should not show ai button when feature toggle is off', () => {
+        createComponent({ aiChatAvailable: true });
+        mockTracking(undefined, wrapper.element, jest.spyOn);
+        expect(findAiAssistnantBtn().exists()).toBe(false);
+      });
+
+      it('should not show ai button when feature is unavailable', () => {
+        createComponent({ aiCiConfigGenerator: true });
+        mockTracking(undefined, wrapper.element, jest.spyOn);
+        expect(findAiAssistnantBtn().exists()).toBe(false);
+      });
+    });
 
+    describe('when feature is available', () => {
+      it('should show ai button', () => {
+        createComponent({ aiCiConfigGenerator: true, aiChatAvailable: true });
+        mockTracking(undefined, wrapper.element, jest.spyOn);
+        expect(findAiAssistnantBtn().exists()).toBe(true);
+      });
+    });
+  });
   describe('link button', () => {
     beforeEach(() => {
       createComponent();
diff --git a/spec/frontend/ci/pipeline_editor/components/pipeline_editor_tabs_spec.js b/spec/frontend/ci/pipeline_editor/components/pipeline_editor_tabs_spec.js
index cbdf01105c7..471b033913b 100644
--- a/spec/frontend/ci/pipeline_editor/components/pipeline_editor_tabs_spec.js
+++ b/spec/frontend/ci/pipeline_editor/components/pipeline_editor_tabs_spec.js
@@ -57,6 +57,7 @@ describe('Pipeline editor tabs component', () => {
         isNewCiConfigFile: true,
         showDrawer: false,
         showJobAssistantDrawer: false,
+        showAiAssistantDrawer: false,
         ...props,
       },
       data() {
@@ -65,6 +66,7 @@ describe('Pipeline editor tabs component', () => {
         };
       },
       provide: {
+        aiChatAvailable: false,
         ciConfigPath: '/path/to/ci-config',
         ciLintPath: mockCiLintPath,
         currentBranch: 'main',