(INITIAL_SELECTED_ENTITIES);
+
+ // Convert the selected values (urns) to EntityAndType format for SearchSelect
+ const selectedEntities = useMemo(
+ () => selectedValues.map((urn) => ({ urn, type: extractTypeFromUrn(urn) })),
+ [selectedValues],
+ );
+
+ // only run this once
+ useEffect(() => {
+ if (tempSelectedEntities === INITIAL_SELECTED_ENTITIES) {
+ setTempSelectedEntities(selectedEntities);
+ }
+ }, [selectedEntities, setTempSelectedEntities, tempSelectedEntities]);
+
+ // call updateSelectedValues when tempSelectedEntities changes
+ useEffect(() => {
+ if (tempSelectedEntities !== INITIAL_SELECTED_ENTITIES) {
+ updateSelectedValues(tempSelectedEntities.map((entity) => entity.urn));
+ }
+ // this is excluded from the dependency array because updateSelectedValues is reconstructed
+ // by the parent component on every render
+ // eslint-disable-next-line react-hooks/exhaustive-deps
+ }, [tempSelectedEntities]);
+
+ const removeEntity = (entity: EntityAndType) => {
+ setTempSelectedEntities(tempSelectedEntities.filter((e) => e.urn !== entity.urn));
+ };
+
+ const hydratedEntityMap = useHydratedEntityMap(tempSelectedEntities.map((entity) => entity.urn));
+
+ return (
+
+
+ Search Values
+
+
+
+
+
+ Selected Values
+
+ {tempSelectedEntities.length === 0 ? (
+ No values selected
+ ) : (
+ tempSelectedEntities.map((entity) => (
+
+ {hydratedEntityMap[entity.urn] ? (
+
+ ) : (
+
+ )}
+
+ removeEntity(entity)} />
+
+
+ ))
+ )}
+
+
+
+ );
+}
diff --git a/datahub-web-react/src/app/entityV2/shared/components/styled/search/tests/SearchSelectUrnInput.test.tsx b/datahub-web-react/src/app/entityV2/shared/components/styled/search/tests/SearchSelectUrnInput.test.tsx
new file mode 100644
index 0000000000..bc6cab2b13
--- /dev/null
+++ b/datahub-web-react/src/app/entityV2/shared/components/styled/search/tests/SearchSelectUrnInput.test.tsx
@@ -0,0 +1,194 @@
+import { MockedProvider } from '@apollo/client/testing';
+import { act, fireEvent, render, screen } from '@testing-library/react';
+import React from 'react';
+import { EntityType } from '../../../../../../../types.generated';
+import TestPageContainer from '../../../../../../../utils/test-utils/TestPageContainer';
+import { EntityAndType } from '../../../../../../entity/shared/types';
+import { SearchSelectUrnInput } from '../SearchSelectUrnInput';
+
+// Mock the useHydratedEntityMap hook
+vi.mock('@src/app/entityV2/shared/tabs/Properties/useHydratedEntityMap', () => ({
+ useHydratedEntityMap: (urns: string[]) => {
+ const map: { [key: string]: any } = {};
+ urns?.forEach((urn) => {
+ // Simple mock entity structure
+ map[urn] = {
+ urn,
+ type: EntityType.Dataset, // Assuming dataset for simplicity
+ properties: { name: `Entity ${urn.split(':').pop()}` },
+ platform: {
+ urn: 'urn:li:dataPlatform:test',
+ name: 'test-platform',
+ properties: {
+ displayName: 'Test Platform',
+ logoUrl: '',
+ },
+ },
+ // Add other necessary fields if EntityLink mock needs them
+ };
+ });
+ return map;
+ },
+}));
+
+// Mock SearchSelect component - capture the setSelectedEntities callback
+const mockSetSelectedEntities = vi.fn(); // Spy will now store the *callback itself* when called
+vi.mock('@src/app/entityV2/shared/components/styled/search/SearchSelect', () => ({
+ SearchSelect: (props: { setSelectedEntities: (entities: EntityAndType[]) => void }) => {
+ // Call the spy with the actual callback function when the mock renders
+ if (props.setSelectedEntities) {
+ // Check if the prop exists before calling
+ mockSetSelectedEntities(props.setSelectedEntities);
+ }
+ return Mock SearchSelect
;
+ },
+}));
+
+// Mock EntityLink component
+vi.mock('@src/app/homeV2/reference/sections/EntityLink', () => ({
+ // eslint-disable-next-line @typescript-eslint/no-explicit-any
+ EntityLink: ({ entity }: { entity: any }) => (
+ {entity.properties?.name || entity.urn}
+ ),
+}));
+
+// Mock Icon component used for removal
+vi.mock('@src/alchemy-components', async () => {
+ // Import the actual module
+ const actual = await vi.importActual('@src/alchemy-components');
+
+ return {
+ // Spread the actual module exports
+ ...actual,
+ // Override specific components with mocks
+ // eslint-disable-next-line @typescript-eslint/no-explicit-any
+ Icon: (props: { icon: string; onClick?: () => void; [key: string]: any }) => (
+
+ ),
+ Tooltip: ({ children }: { children: React.ReactNode }) => <>{children},
+ Button: ({ children, onClick }: { children: React.ReactNode; onClick?: () => void }) => (
+
+ ),
+ // No need to mock colors and typography anymore, we use the actual ones
+ };
+});
+
+describe('SearchSelectUrnInput', () => {
+ const mockUpdateSelectedValues = vi.fn();
+
+ beforeEach(() => {
+ mockUpdateSelectedValues.mockClear();
+ mockSetSelectedEntities.mockClear();
+ });
+
+ it('renders correctly with no initial selection', () => {
+ render(
+
+
+
+
+ ,
+ );
+
+ expect(screen.getByText('Search Values')).toBeInTheDocument();
+ expect(screen.getByTestId('mock-search-select')).toBeInTheDocument();
+ expect(screen.getByText('Selected Values')).toBeInTheDocument();
+ expect(screen.getByText('No values selected')).toBeInTheDocument();
+ });
+
+ it('renders with initial selected values', () => {
+ const initialUrns = ['urn:li:dataset:1', 'urn:li:dataset:2'];
+ render(
+
+
+
+
+ ,
+ );
+
+ expect(screen.getByText('Selected Values')).toBeInTheDocument();
+ expect(screen.getByTestId('entity-link-urn:li:dataset:1')).toHaveTextContent('Entity 1');
+ expect(screen.getByTestId('entity-link-urn:li:dataset:2')).toHaveTextContent('Entity 2');
+ expect(screen.queryByText('No values selected')).not.toBeInTheDocument();
+ });
+
+ it('calls updateSelectedValues when a value is removed', () => {
+ const initialUrns = ['urn:li:dataset:1', 'urn:li:dataset:2'];
+ render(
+
+
+
+
+ ,
+ );
+
+ // Find the remove button for the first entity
+ const removeButton = screen.getAllByTestId('mock-icon-x')[0]; // Assuming order matches initialUrns
+ fireEvent.click(removeButton);
+
+ // Check if updateSelectedValues was called with the remaining URN
+ // It gets called twice: once on initial load (useEffect) and once on removal
+ expect(mockUpdateSelectedValues).toHaveBeenCalledTimes(2);
+ expect(mockUpdateSelectedValues).toHaveBeenLastCalledWith(['urn:li:dataset:2']);
+
+ // Check if the removed entity is no longer rendered
+ expect(screen.queryByTestId('entity-link-urn:li:dataset:1')).not.toBeInTheDocument();
+ expect(screen.getByTestId('entity-link-urn:li:dataset:2')).toBeInTheDocument();
+ });
+
+ it('calls updateSelectedValues when SearchSelect updates selection', () => {
+ const initialUrns: string[] = [];
+ const newEntity: EntityAndType = { urn: 'urn:li:dataset:3', type: EntityType.Dataset };
+ render(
+
+
+
+
+ ,
+ );
+
+ // Ensure the SearchSelect mock rendered and called our spy with the callback
+ expect(mockSetSelectedEntities).toHaveBeenCalled();
+
+ // Get the actual callback function from the spy's arguments
+ const capturedSetSelectedEntities = mockSetSelectedEntities.mock.calls[0][0];
+ expect(capturedSetSelectedEntities).toBeInstanceOf(Function); // Verify it's a function
+
+ // Simulate SearchSelect calling its setSelectedEntities prop within act
+ act(() => {
+ capturedSetSelectedEntities([newEntity]);
+ });
+
+ // Check if updateSelectedValues was called with the new URN
+ // It gets called once on initial load (useEffect) and once when the selection changes
+ expect(mockUpdateSelectedValues).toHaveBeenCalledTimes(2);
+ expect(mockUpdateSelectedValues).toHaveBeenLastCalledWith([newEntity.urn]);
+
+ // Check if the new entity is rendered
+ expect(screen.getByTestId('entity-link-urn:li:dataset:3')).toHaveTextContent('Entity 3');
+ });
+});