diff --git a/e2e-tests/cypress/e2e/browse-data.cy.js b/e2e-tests/cypress/e2e/browse-data.cy.js
index 51a0127cb..b1dd797ef 100644
--- a/e2e-tests/cypress/e2e/browse-data.cy.js
+++ b/e2e-tests/cypress/e2e/browse-data.cy.js
@@ -349,4 +349,17 @@ describe('Data browser data', () => {
     cy.testid('CompareModelTab_tabOperations').click();
     cy.themeshot('comparesettings');
   });
+
+  it.only('Query editor - AI assistant', () => {
+    cy.contains('MySql-connection').click();
+    cy.contains('MyChinook').click();
+    cy.testid('TabsPanel_buttonNewQuery').click();
+    cy.testid('QueryTab_switchAiAssistantButton').click();
+    cy.testid('QueryAiAssistant_promptInput').type('album names');
+    cy.testid('QueryAiAssistant_queryFromQuestionButton').click();
+    cy.contains('Use this').click();
+    cy.testid('QueryTab_executeButton').click();
+    cy.contains('Balls to the Wall');
+    cy.themeshot('aiassistant');
+  });
 });
diff --git a/packages/web/src/buttons/ToolStripCommandSplitButton.svelte b/packages/web/src/buttons/ToolStripCommandSplitButton.svelte
index 933b5fe51..0373bb8cc 100644
--- a/packages/web/src/buttons/ToolStripCommandSplitButton.svelte
+++ b/packages/web/src/buttons/ToolStripCommandSplitButton.svelte
@@ -17,4 +17,5 @@
   {menu}
   {hideDisabled}
   {buttonLabel}
+  {...$$restProps}
 />
diff --git a/packages/web/src/buttons/ToolStripSplitButton.svelte b/packages/web/src/buttons/ToolStripSplitButton.svelte
index 007ca44cd..a21253c85 100644
--- a/packages/web/src/buttons/ToolStripSplitButton.svelte
+++ b/packages/web/src/buttons/ToolStripSplitButton.svelte
@@ -21,7 +21,7 @@
 
 <div class="button" class:disabled {title}>
   <div class="inner" class:disabled>
-    <div class="main" class:disabled on:click={handleClick}>
+    <div class="main" class:disabled on:click={handleClick} data-testid={$$props['data-testid']}>
       <span class="icon" class:disabled><FontIcon {icon} /></span>
       <slot />
     </div>
diff --git a/packages/web/src/buttons/ToolStripSplitDropDownButton.svelte b/packages/web/src/buttons/ToolStripSplitDropDownButton.svelte
index ffec46466..5236004f7 100644
--- a/packages/web/src/buttons/ToolStripSplitDropDownButton.svelte
+++ b/packages/web/src/buttons/ToolStripSplitDropDownButton.svelte
@@ -19,6 +19,6 @@
   }
 </script>
 
-<ToolStripSplitButton {title} {icon} {disabled} on:splitclick={handleClick} on:click>
+<ToolStripSplitButton {title} {icon} {disabled} on:splitclick={handleClick} on:click {...$$restProps}>
   <slot />
 </ToolStripSplitButton>
diff --git a/packages/web/src/modals/ConfirmModal.svelte b/packages/web/src/modals/ConfirmModal.svelte
index 56cafc438..dceb19281 100644
--- a/packages/web/src/modals/ConfirmModal.svelte
+++ b/packages/web/src/modals/ConfirmModal.svelte
@@ -21,8 +21,14 @@
           closeCurrentModal();
           onConfirm();
         }}
+        data-testid="ConfirmModal_okButton"
+      />
+      <FormStyledButton
+        type="button"
+        value="Close"
+        on:click={closeCurrentModal}
+        data-testid="ConfirmModal_closeButton"
       />
-      <FormStyledButton type="button" value="Close" on:click={closeCurrentModal} />
     </svelte:fragment>
   </ModalBase>
 </FormProvider>
diff --git a/packages/web/src/tabs/QueryTab.svelte b/packages/web/src/tabs/QueryTab.svelte
index 608244df6..9113ccd5c 100644
--- a/packages/web/src/tabs/QueryTab.svelte
+++ b/packages/web/src/tabs/QueryTab.svelte
@@ -538,12 +538,16 @@
             editor.setValue(text);
           }
         }}
+        {tabid}
       />
     </svelte:fragment>
   </HorizontalSplitter>
   <svelte:fragment slot="toolstrip">
-    <ToolStripCommandSplitButton commands={['query.execute', 'query.executeCurrent']} />
-    <ToolStripCommandButton command="query.kill" />
+    <ToolStripCommandSplitButton
+      commands={['query.execute', 'query.executeCurrent']}
+      data-testid="QueryTab_executeButton"
+    />
+    <ToolStripCommandButton command="query.kill" data-testid="QueryTab_killButton" />
     <ToolStripSaveButton idPrefix="query" />
     <ToolStripCommandButton command="query.formatCode" />
     {#if resultCount == 1}
@@ -562,7 +566,13 @@
       icon="icon at"
       title="Query parameter style"
     />
-    <ToolStripCommandButton command="query.switchAiAssistant" hideDisabled />
+    <ToolStripCommandButton
+      command="query.switchAiAssistant"
+      hideDisabled
+      data-testid="QueryTab_switchAiAssistantButton"
+    >
+      AI Assistant
+    </ToolStripCommandButton>
   </svelte:fragment>
 </ToolStripContainer>