test: fix add and remove test

Author: gagik

packages/compass-components/src/components/document-list/element.spec.tsx

@@ -3,7 +3,11 @@ import { render, screen, userEvent } from '@mongodb-js/testing-library-compass';
 import { expect } from 'chai';
 import sinon from 'sinon';
 import HadronDocument from 'hadron-document';
-import { HadronElement, getNestedKeyPathForElement } from './element';
+import {
+  HadronElement,
+  type QueryBarController,
+  getNestedKeyPathForElement,
+} from './element';
 import type { Element } from 'hadron-document';
 
 describe('HadronElement', function () {
@@ -32,21 +36,18 @@ describe('HadronElement', function () {
 
       // Mock queryBar controller
       const mockQueryBar = {
-        isInQuery: sinon.stub(),
+        isInQuery: sinon.stub().returns(false),
         toggleQueryFilter: sinon.spy(),
       };
 
-      // Initially not in query
-      mockQueryBar.isInQuery.returns(false);
-
       const { rerender } = render(
         <HadronElement
           value={nestedElement}
           editable={true}
           editingEnabled={true}
           lineNumberSize={1}
           onAddElement={() => {}}
-          queryBar={mockQueryBar}
+          queryBar={mockQueryBar as unknown as QueryBarController}
         />
       );
 
@@ -57,12 +58,10 @@ describe('HadronElement', function () {
       expect(screen.getByText('Add to query')).to.exist;
       expect(screen.queryByText('Remove from query')).to.not.exist;
 
-      // Click "Add to query"
       userEvent.click(screen.getByText('Add to query'), undefined, {
         skipPointerEventsCheck: true,
       });
 
-      // Verify toggleQueryFilter was called with correct parameters
       expect(mockQueryBar.toggleQueryFilter).to.have.been.calledWith(
         'user.name',
         nestedElement.generateObject()
@@ -79,7 +78,10 @@ describe('HadronElement', function () {
           editingEnabled={true}
           lineNumberSize={1}
           onAddElement={() => {}}
-          queryBar={mockQueryBar}
+          queryBar={{
+            ...mockQueryBar,
+            isInQuery: sinon.stub().returns(true),
+          }}
         />
       );
 
@@ -89,12 +91,10 @@ describe('HadronElement', function () {
       expect(screen.getByText('Remove from query')).to.exist;
       expect(screen.queryByText('Add to query')).to.not.exist;
 
-      // Click "Remove from query"
       userEvent.click(screen.getByText('Remove from query'), undefined, {
         skipPointerEventsCheck: true,
       });
 
-      // Verify toggleQueryFilter was called again
       expect(mockQueryBar.toggleQueryFilter).to.have.been.calledTwice;
       expect(mockQueryBar.toggleQueryFilter.secondCall).to.have.been.calledWith(
         'user.name',

packages/compass-components/src/components/document-list/element.tsx

@@ -504,20 +504,22 @@ export const HadronElement: React.FunctionComponent<{
     collapse,
   } = useHadronElement(element);
 
+  const { isInQuery, toggleQueryFilter } = queryBar || {};
+
   // Add context menu hook for the field
   const fieldContextMenuRef = useContextMenuItems(
     () => [
-      ...(queryBar
+      ...(isInQuery && toggleQueryFilter
         ? [
             {
-              label: queryBar.isInQuery(
+              label: isInQuery(
                 getNestedKeyPathForElement(element),
                 element.generateObject()
               )
                 ? 'Remove from query'
                 : 'Add to query',
               onAction: () => {
-                queryBar.toggleQueryFilter(
+                toggleQueryFilter(
                   getNestedKeyPathForElement(element),
                   element.generateObject()
                 );
@@ -544,7 +546,7 @@ export const HadronElement: React.FunctionComponent<{
           ]
         : []),
     ],
-    [element, key.value, value.value, type.value, queryBar]
+    [element, key.value, value.value, type.value, isInQuery, toggleQueryFilter]
   );
 
   const toggleExpanded = () => {