@finos/legend-application-studio 15.1.0 → 15.2.0

package/src/components/editor/edit-panel/connection-editor/post-processor-editor/MapperPostProcessorEditor.tsx

@@ -0,0 +1,278 @@
+/**
+ * Copyright (c) 2020-present, Goldman Sachs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import {
+  ResizablePanelGroup,
+  ResizablePanel,
+  PanelHeaderActionItem,
+  PlusIcon,
+  ResizablePanelSplitter,
+  PanelTextEditor,
+  PanelForm,
+  ContextMenu,
+  MenuContent,
+  MenuContentItem,
+  DropdownMenu,
+  PanelListSelectorItem,
+  BlankPanelContent,
+  ResizablePanelSplitterLine,
+  Panel,
+  PanelHeader,
+  PanelContent,
+} from '@finos/legend-art';
+import {
+  type Mapper,
+  MapperPostProcessor,
+  TableNameMapper,
+  SchemaNameMapper,
+} from '@finos/legend-graph';
+import { observer } from 'mobx-react-lite';
+import type { MapperPostProcessorEditorState } from '../../../../../stores/editor-state/element-editor-state/connection/PostProcessorEditorState.js';
+import {
+  mapperPostProcessor_addMapper,
+  mapperPostProcessor_deleteMapper,
+  mapper_setFrom,
+  mapper_setTo,
+  schemaNameMapper_setFrom,
+  schemaNameMapper_setTo,
+} from '../../../../../stores/graphModifier/StoreRelational_GraphModifierHelper.js';
+
+export const MapperPostProcessorEditor = observer(
+  (props: {
+    postProcessor: MapperPostProcessor;
+    isReadOnly: boolean;
+    postProcessorState: MapperPostProcessorEditorState;
+  }) => {
+    const { postProcessorState, isReadOnly, postProcessor } = props;
+    const selectedMapper = postProcessorState.selectedMapper;
+
+    const selectedSchemaNameMapper =
+      selectedMapper instanceof TableNameMapper
+        ? selectedMapper.schema
+        : undefined;
+
+    const selectMapper = (val: Mapper): void => {
+      postProcessorState.setSelectedMapper(val);
+    };
+
+    const addSchemaMapper = (): void => {
+      mapperPostProcessor_addMapper(
+        postProcessor,
+        new SchemaNameMapper('', ''),
+      );
+      postProcessorState.setSelectedMapper(postProcessor.mappers.at(-1));
+    };
+
+    const addTableMapper = (): void => {
+      mapperPostProcessor_addMapper(
+        postProcessor,
+        new TableNameMapper('', '', new SchemaNameMapper('', '')),
+      );
+      postProcessorState.setSelectedMapper(postProcessor.mappers.at(-1));
+    };
+
+    const deleteMapper =
+      (mapper: Mapper): (() => void) =>
+      (): void => {
+        if (postProcessorState.postProcessor instanceof MapperPostProcessor) {
+          mapperPostProcessor_deleteMapper(
+            postProcessorState.postProcessor,
+            mapper,
+          );
+        }
+
+        if (mapper === postProcessorState.selectedMapper) {
+          postProcessorState.setSelectedMapper(undefined);
+        }
+      };
+
+    const mappers = postProcessor.mappers;
+
+    const isSchemaMapperDuplicated = (val: Mapper): boolean =>
+      mappers.filter(
+        (mapper) =>
+          mapper.from === val.from &&
+          mapper.to === val.to &&
+          val instanceof SchemaNameMapper &&
+          mapper instanceof SchemaNameMapper,
+      ).length >= 2;
+
+    const isTableMapperDuplicated = (val: Mapper): boolean =>
+      mappers
+        .filter(
+          (mapper) =>
+            val instanceof TableNameMapper && mapper instanceof TableNameMapper,
+        )
+        .filter(
+          (mapper) =>
+            (mapper as TableNameMapper).schema.from ===
+              (val as TableNameMapper).schema.from &&
+            (mapper as TableNameMapper).schema.from ===
+              (val as TableNameMapper).schema.from &&
+            (mapper as TableNameMapper).from ===
+              (val as TableNameMapper).from &&
+            (mapper as TableNameMapper).from === (val as TableNameMapper).from,
+        ).length >= 2;
+
+    const isMapperDuplicated = (val: Mapper): boolean =>
+      isSchemaMapperDuplicated(val) || isTableMapperDuplicated(val);
+
+    return (
+      <ResizablePanelGroup orientation="horizontal">
+        <ResizablePanel>
+          <ResizablePanelGroup orientation="vertical">
+            <ResizablePanel size={170} minSize={70}>
+              <Panel>
+                <PanelHeader title="mapper">
+                  <DropdownMenu
+                    disabled={isReadOnly}
+                    content={
+                      <MenuContent>
+                        <MenuContentItem onClick={addSchemaMapper}>
+                          New Schema Mapper
+                        </MenuContentItem>
+                        <MenuContentItem onClick={addTableMapper}>
+                          New Table Mapper
+                        </MenuContentItem>
+                      </MenuContent>
+                    }
+                    menuProps={{
+                      anchorOrigin: { vertical: 'bottom', horizontal: 'right' },
+                      transformOrigin: { vertical: 'top', horizontal: 'right' },
+                      elevation: 7,
+                    }}
+                  >
+                    <PanelHeaderActionItem
+                      disabled={isReadOnly}
+                      title="Create Mapper"
+                    >
+                      <PlusIcon />
+                    </PanelHeaderActionItem>
+                  </DropdownMenu>
+                </PanelHeader>
+                <PanelContent>
+                  {mappers.map((mapper) => (
+                    <ContextMenu
+                      key={mapper._UUID}
+                      disabled={isReadOnly}
+                      content={
+                        <MenuContent>
+                          <MenuContentItem
+                            disabled={isReadOnly}
+                            onClick={deleteMapper(mapper)}
+                          >
+                            Delete
+                          </MenuContentItem>
+                        </MenuContent>
+                      }
+                      menuProps={{ elevation: 7 }}
+                    >
+                      <PanelListSelectorItem
+                        title={
+                          mapper instanceof TableNameMapper
+                            ? 'Table Mapper'
+                            : 'Schema Mapper'
+                        }
+                        validationErrorMessage={
+                          isMapperDuplicated(mapper)
+                            ? 'Mappers have the same values'
+                            : undefined
+                        }
+                        isSelected={
+                          mapper === postProcessorState.selectedMapper
+                        }
+                        onSelect={() => selectMapper(mapper)}
+                      />
+                    </ContextMenu>
+                  ))}
+                </PanelContent>{' '}
+              </Panel>
+            </ResizablePanel>
+            <ResizablePanelSplitter>
+              <ResizablePanelSplitterLine color="var(--color-dark-grey-200)" />
+            </ResizablePanelSplitter>
+            <ResizablePanel>
+              <Panel>
+                <PanelHeader
+                  title={
+                    postProcessorState.selectedMapper instanceof TableNameMapper
+                      ? 'Table Mapper'
+                      : 'Schema Mapper'
+                  }
+                ></PanelHeader>
+                <PanelContent>
+                  {!selectedMapper && (
+                    <BlankPanelContent>
+                      {!mappers.length
+                        ? 'Add a mapper'
+                        : 'Select a mapper to view'}
+                    </BlankPanelContent>
+                  )}
+                  {selectedMapper && (
+                    <PanelForm>
+                      <PanelTextEditor
+                        isReadOnly={isReadOnly}
+                        value={selectedMapper.from}
+                        name="From"
+                        update={(value: string | undefined): void =>
+                          mapper_setFrom(selectedMapper, value ?? '')
+                        }
+                      />
+                      <PanelTextEditor
+                        isReadOnly={isReadOnly}
+                        value={selectedMapper.to}
+                        name="To"
+                        update={(value: string | undefined): void =>
+                          mapper_setTo(selectedMapper, value ?? '')
+                        }
+                      />
+                      {selectedSchemaNameMapper && (
+                        <>
+                          <PanelTextEditor
+                            isReadOnly={isReadOnly}
+                            value={selectedSchemaNameMapper.from}
+                            name="Schema - From"
+                            update={(value: string | undefined): void => {
+                              schemaNameMapper_setFrom(
+                                selectedSchemaNameMapper,
+                                value ?? '',
+                              );
+                            }}
+                          />
+                          <PanelTextEditor
+                            isReadOnly={isReadOnly}
+                            value={selectedSchemaNameMapper.to}
+                            name="Schema - To"
+                            update={(value: string | undefined): void =>
+                              schemaNameMapper_setTo(
+                                selectedSchemaNameMapper,
+                                value ?? '',
+                              )
+                            }
+                          />
+                        </>
+                      )}
+                    </PanelForm>
+                  )}
+                </PanelContent>
+              </Panel>
+            </ResizablePanel>
+          </ResizablePanelGroup>
+        </ResizablePanel>
+      </ResizablePanelGroup>
+    );
+  },
+);

package/src/components/editor/edit-panel/data-editor/EmbeddedDataEditor.tsx

@@ -157,7 +157,8 @@ export const DataElementReferenceDataEditor = observer(
     const dataElement =
       dataElementReferenceState.embeddedData.dataElement.value;
     const editorStore = dataElementReferenceState.editorStore;
-    const options = editorStore.dataOptions;
+    const options =
+      editorStore.graphManagerState.usableDataElements.map(buildElementOption);
     const selectedOption = buildElementOption(dataElement);
     const onDataElementChange = (val: {
       label: string;

package/src/components/editor/edit-panel/mapping-editor/EnumerationMappingEditor.tsx

@@ -89,7 +89,11 @@ const EnumerationMappingSourceSelectorModal = observer(
       ),
     ]
       .map(buildElementOption)
-      .concat(editorStore.enumerationOptions);
+      .concat(
+        editorStore.graphManagerState.usableEnumerations.map(
+          buildElementOption,
+        ),
+      );
 
     const sourceSelectorRef = useRef<SelectComponent>(null);
     const filterOption = createFilter({

package/src/components/editor/edit-panel/mapping-editor/InstanceSetImplementationEditor.tsx

@@ -363,7 +363,7 @@ export const InstanceSetImplementationSourceExplorer = observer(
             {isUnsupported && (
               <UnsupportedEditorPanel
                 isReadOnly={isReadOnly}
-                text={`Can't display class mapping source in form mode`}
+                text="Can't display class mapping source in form mode"
               ></UnsupportedEditorPanel>
             )}
             {sourceElementForSourceSelectorModal !== undefined && (
@@ -573,7 +573,7 @@ export const InstanceSetImplementationEditor = observer(
                     {isUnsupported && (
                       <UnsupportedEditorPanel
                         isReadOnly={isReadOnly}
-                        text={`Can't display class mapping in form mode`}
+                        text="Can't display class mapping in form mode"
                       ></UnsupportedEditorPanel>
                     )}
                   </div>

package/src/components/editor/edit-panel/mapping-editor/InstanceSetImplementationSourceSelectorModal.tsx

@@ -128,22 +128,15 @@ export const InstanceSetImplementationSourceSelectorModal = observer(
     const editorStore = useEditorStore();
     const applicationStore = useApplicationStore();
     const options = (
-      editorStore.graphManagerState.graph.ownClasses as MappingElementSource[]
+      editorStore.graphManagerState.usableClasses as MappingElementSource[]
     )
-      .concat(
-        editorStore.graphManagerState.graph.dependencyManager
-          .classes as MappingElementSource[],
-      )
       .concat(
         editorStore.graphManagerState.graph.ownFlatDatas.flatMap(
           getAllRecordTypes,
         ),
       )
       .concat(
-        editorStore.graphManagerState.graph.ownDatabases
-          .concat(
-            editorStore.graphManagerState.graph.dependencyManager.databases,
-          )
+        editorStore.graphManagerState.usableDatabases
           .flatMap((e) =>
             e.schemas.flatMap((schema) =>
               (schema.tables as (Table | View)[]).concat(schema.views),

package/src/components/editor/edit-panel/mapping-editor/NewMappingElementModal.tsx

@@ -45,6 +45,7 @@ import {
 } from '@finos/legend-graph';
 import { BASIC_SET_IMPLEMENTATION_TYPE } from '../../../../stores/shared/ModelUtil.js';
 import {
+  buildElementOption,
   getPackageableElementOptionFormatter,
   type PackageableElementOption,
 } from '@finos/legend-application';
@@ -77,9 +78,9 @@ export const NewMappingElementModal = observer(() => {
   // Target
   const targetSelectorRef = useRef<SelectComponent>(null);
   const options: PackageableElementOption<PackageableElement>[] = [
-    ...editorStore.enumerationOptions,
-    ...editorStore.associationOptions,
-    ...editorStore.classOptions,
+    ...editorStore.graphManagerState.usableEnumerations.map(buildElementOption),
+    ...editorStore.graphManagerState.usableAssociations.map(buildElementOption),
+    ...editorStore.graphManagerState.usableClasses.map(buildElementOption),
   ].sort(compareLabelFn);
   const filterOption = createFilter({
     ignoreCase: true,

package/src/components/editor/edit-panel/mapping-editor/relational/TableOrViewSourceTree.tsx

@@ -35,7 +35,7 @@ import {
 import {
   generateColumnTypeLabel,
   renderColumnTypeIcon,
-} from '../../../../../stores/editor-state/element-editor-state/mapping/relational/DatabaseEditorHelper.js';
+} from '../../connection-editor/DatabaseEditorHelper.js';
 import {
   type Type,
   type Table,

package/src/components/editor/edit-panel/service-editor/ServiceExecutionEditor.tsx

@@ -68,7 +68,10 @@ import {
   validate_PureExecutionMapping,
 } from '@finos/legend-graph';
 import { guaranteeNonNullable } from '@finos/legend-shared';
-import type { PackageableElementOption } from '@finos/legend-application';
+import {
+  buildElementOption,
+  type PackageableElementOption,
+} from '@finos/legend-application';
 
 const PureExecutionContextConfigurationEditor = observer(
   (props: {
@@ -87,7 +90,8 @@ const PureExecutionContextConfigurationEditor = observer(
       executionContext.mapping.value,
     );
     const mapping = executionContext.mapping.value;
-    const mappingOptions = editorStore.mappingOptions;
+    const mappingOptions =
+      editorStore.graphManagerState.usableMappings.map(buildElementOption);
     const noMappingLabel = (
       <div
         className="service-execution-editor__configuration__mapping-option--empty"

package/src/components/editor/edit-panel/service-editor/testable/ServiceTestDataEditor.tsx

@@ -46,13 +46,18 @@ import type { EmbeddedDataTypeOption } from '../../../../../stores/editor-state/
 import { EmbeddedDataEditor } from '../../data-editor/EmbeddedDataEditor.js';
 import { EmbeddedDataType } from '../../../../../stores/editor-state/ExternalFormatState.js';
 import { flowResult } from 'mobx';
-import { buildElementOption } from '@finos/legend-application';
+import {
+  buildElementOption,
+  useApplicationStore,
+} from '@finos/legend-application';
 import { prettyCONSTName } from '@finos/legend-shared';
 import type { DSLData_LegendStudioApplicationPlugin_Extension } from '../../../../../stores/DSLData_LegendStudioApplicationPlugin_Extension.js';
+import { useEditorStore } from '../../../EditorStoreProvider.js';
 
 export const ConnectionTestDataEditor = observer(
   (props: { connectionTestDataState: ConnectionTestDataState }) => {
     const { connectionTestDataState } = props;
+    const applicationStore = useApplicationStore();
     const isReadOnly =
       connectionTestDataState.testDataState.testSuiteState.testableState
         .serviceEditorState.isReadOnly;
@@ -66,10 +71,10 @@ export const ConnectionTestDataEditor = observer(
     };
     const generateTestData = (): void => {
       flowResult(connectionTestDataState.generateTestData()).catch(
-        connectionTestDataState.editorStore.applicationStore
-          .alertUnhandledError,
+        applicationStore.alertUnhandledError,
       );
     };
+
     return (
       <div className="service-test-data-editor">
         <div className="service-test-suite-editor__header">
@@ -209,7 +214,9 @@ const ConnectionTestDataItem = observer(
 export const NewConnectionDataModal = observer(
   (props: { testDataState: ServiceTestDataState }) => {
     const { testDataState } = props;
-    const dataElementOptions = testDataState.editorStore.dataOptions;
+    const editorStore = useEditorStore();
+    const dataElementOptions =
+      editorStore.graphManagerState.usableDataElements.map(buildElementOption);
     const newConnectionState = testDataState.newConnectionDataState;
     const dataElement = newConnectionState.dataElement;
     const selectedDataElement = dataElement
@@ -271,7 +278,7 @@ export const NewConnectionDataModal = observer(
           value: newConnectionState.embeddedDataType.value,
         }
       : undefined;
-    const extraOptionTypes = testDataState.editorStore.pluginManager
+    const extraOptionTypes = editorStore.pluginManager
       .getApplicationPlugins()
       .flatMap(
         (plugin) =>

package/src/components/editor/edit-panel/uml-editor/AssociationEditor.tsx

@@ -89,6 +89,7 @@ import {
   getClassPropertyType,
 } from '../../../../stores/shared/ModelUtil.js';
 import {
+  buildElementOption,
   useApplicationNavigationContext,
   type PackageableElementOption,
 } from '@finos/legend-application';
@@ -119,7 +120,8 @@ const AssociationPropertyBasicEditor = observer(
     // Generic Type
     const [isEditingType, setIsEditingType] = useState(false);
     // TODO: make this so that association can only refer to classes from the same graph space
-    const propertyTypeOptions = editorStore.classOptions;
+    const propertyTypeOptions =
+      editorStore.graphManagerState.usableClasses.map(buildElementOption);
     const propertyType = property.genericType.value.rawType;
     const propertyTypeName = getClassPropertyType(propertyType);
     const filterOption = createFilter({

package/src/components/editor/edit-panel/uml-editor/ClassEditor.tsx

@@ -96,6 +96,7 @@ import {
 import { StudioLambdaEditor } from '../../../shared/StudioLambdaEditor.js';
 import {
   ApplicationNavigationContextData,
+  buildElementOption,
   getPackageableElementOptionFormatter,
   useApplicationNavigationContext,
   useApplicationStore,
@@ -167,7 +168,10 @@ const PropertyBasicEditor = observer(
     };
     // Generic Type
     const [isEditingType, setIsEditingType] = useState(false);
-    const propertyTypeOptions = editorStore.classPropertyGenericTypeOptions;
+    const propertyTypeOptions =
+      editorStore.graphManagerState.usableClassPropertyTypes.map(
+        buildElementOption,
+      );
     const propertyType = property.genericType.value.rawType;
     const propertyTypeName = getClassPropertyType(propertyType);
     const filterOption = createFilter({
@@ -492,7 +496,10 @@ const DerivedPropertyBasicEditor = observer(
       property_setName(derivedProperty, event.target.value);
     // Generic Type
     const [isEditingType, setIsEditingType] = useState(false);
-    const propertyTypeOptions = editorStore.classPropertyGenericTypeOptions;
+    const propertyTypeOptions =
+      editorStore.graphManagerState.usableClassPropertyTypes.map(
+        buildElementOption,
+      );
     const propertyType = derivedProperty.genericType.value.rawType;
     const propertyTypeName = getClassPropertyType(propertyType);
     const filterOption = createFilter({
@@ -998,16 +1005,18 @@ const SuperTypeEditor = observer(
     const { superType, _class, deleteSuperType, isReadOnly } = props;
     const editorStore = useEditorStore();
     // Type
-    const superTypeOptions = editorStore.classOptions.filter(
-      (classOption) =>
-        classOption.value instanceof Class &&
-        // Exclude current class
-        classOption.value !== _class &&
-        // Exclude super types of the class
-        !getAllSuperclasses(_class).includes(classOption.value) &&
-        // Ensure there is no loop (might be expensive)
-        !getAllSuperclasses(classOption.value).includes(_class),
-    );
+    const superTypeOptions = editorStore.graphManagerState.usableClasses
+      .filter(
+        (c) =>
+          c instanceof Class &&
+          // Exclude current class
+          c !== _class &&
+          // Exclude super types of the class
+          !getAllSuperclasses(_class).includes(c) &&
+          // Ensure there is no loop (might be expensive)
+          !getAllSuperclasses(c).includes(_class),
+      )
+      .map(buildElementOption);
 
     // Drag and Drop
     const handleHover = useCallback(

package/src/components/editor/edit-panel/uml-editor/StereotypeSelector.tsx

@@ -38,7 +38,10 @@ import {
   annotatedElement_swapStereotypes,
   stereotypeReference_setValue,
 } from '../../../../stores/graphModifier/DomainGraphModifierHelper.js';
-import type { PackageableElementOption } from '@finos/legend-application';
+import {
+  buildElementOption,
+  type PackageableElementOption,
+} from '@finos/legend-application';
 import { useDrop, useDrag } from 'react-dnd';
 
 interface StereotypeOption {
@@ -81,9 +84,9 @@ export const StereotypeSelector = observer(
     const editorStore = useEditorStore();
 
     // Profile
-    const profileOptions = editorStore.profileOptions.filter(
-      (p) => p.value.p_stereotypes.length,
-    );
+    const profileOptions = editorStore.graphManagerState.usableProfiles
+      .map(buildElementOption)
+      .filter((p) => p.value.p_stereotypes.length);
     const filterOption = createFilter({
       ignoreCase: true,
       ignoreAccents: false,

package/src/components/editor/edit-panel/uml-editor/TaggedValueEditor.tsx

@@ -42,7 +42,10 @@ import {
   taggedValue_setTag,
   annotatedElement_swapTaggedValues,
 } from '../../../../stores/graphModifier/DomainGraphModifierHelper.js';
-import type { PackageableElementOption } from '@finos/legend-application';
+import {
+  buildElementOption,
+  type PackageableElementOption,
+} from '@finos/legend-application';
 import { useDrop, useDrag } from 'react-dnd';
 
 interface TagOption {
@@ -88,9 +91,9 @@ export const TaggedValueEditor = observer(
       HTMLTextAreaElement | HTMLInputElement
     > = (event) => taggedValue_setValue(taggedValue, event.target.value);
     // Profile
-    const profileOptions = editorStore.profileOptions.filter(
-      (p) => p.value.p_tags.length,
-    );
+    const profileOptions = editorStore.graphManagerState.usableProfiles
+      .map(buildElementOption)
+      .filter((p) => p.value.p_tags.length);
     const profileFilterOption = createFilter({
       ignoreCase: true,
       ignoreAccents: false,

package/src/components/editor/side-bar/CreateNewElementModal.tsx

@@ -41,6 +41,7 @@ import {
 import type { FileGenerationTypeOption } from '../../../stores/editor-state/GraphGenerationState.js';
 import { flowResult } from 'mobx';
 import {
+  buildElementOption,
   getPackageableElementOptionFormatter,
   useApplicationStore,
   type PackageableElementOption,
@@ -164,7 +165,8 @@ const NewRuntimeDriverEditor = observer(() => {
   );
   // mapping
   const mapping = newRuntimeDriver.mapping;
-  const mappingOptions = editorStore.mappingOptions;
+  const mappingOptions =
+    editorStore.graphManagerState.usableMappings.map(buildElementOption);
   const selectedMappingOption = { label: mapping?.path ?? '', value: mapping };
   const onMappingSelectionChange = (
     val: PackageableElementOption<Mapping>,
@@ -216,7 +218,10 @@ const NewPureModelConnectionDriverEditor = observer(
     }): void => newConnectionDriver.setStore(val.value);
     // class
     const _class = newConnectionValueDriver.class;
-    const classOptions = editorStore.classOptions.slice().sort(compareLabelFn);
+    const classOptions = editorStore.graphManagerState.usableClasses
+      .map(buildElementOption)
+      .slice()
+      .sort(compareLabelFn);
     const selectedClassOption = _class
       ? { label: _class.path, value: _class }
       : null;
@@ -321,7 +326,8 @@ const NewServiceDriverEditor = observer(() => {
     editorStore.newElementState.getNewElementDriver(NewServiceDriver);
   // mapping
   const currentMappingOption = newServiceDriver.mappingOption;
-  const mappingOptions = editorStore.mappingOptions;
+  const mappingOptions =
+    editorStore.graphManagerState.usableMappings.map(buildElementOption);
   const onMappingChange = (
     val: PackageableElementOption<Mapping> | null,
   ): void => {

package/src/components/editor/side-bar/testable/GlobalTestRunner.tsx

@@ -38,6 +38,7 @@ import {
 import {
   AssertFail,
   EqualToJsonAssertFail,
+  PackageableElement,
   TestError,
 } from '@finos/legend-graph';
 import { type GeneratorFn, isNonNullable } from '@finos/legend-shared';
@@ -58,6 +59,7 @@ import {
 } from '../../../../stores/sidebar-state/testable/GlobalTestRunnerState.js';
 import { LEGEND_STUDIO_TEST_ID } from '../../../LegendStudioTestID.js';
 import { TextDiffView } from '../../../shared/DiffView.js';
+import { getElementTypeIcon } from '../../../shared/ElementIconUtils.js';
 import { StudioTextInputEditor } from '../../../shared/StudioTextInputEditor.js';
 import { useEditorStore } from '../../EditorStoreProvider.js';
 
@@ -249,6 +251,7 @@ const TestableTreeNodeContainer: React.FC<
 > = (props) => {
   const { node, level, stepPaddingInRem, onNodeSelect } = props;
   const { treeData, testableState, globalTestRunnerState } = props.innerProps;
+  const editorStore = useEditorStore();
   const results = testableState.results;
   const expandIcon =
     node instanceof AssertionTestTreeNodeData ? (
@@ -259,7 +262,16 @@ const TestableTreeNodeContainer: React.FC<
       <ChevronRightIcon />
     );
   const nodeIcon =
-    node instanceof TestableTreeNodeData ? node.testableMetadata.icon : null;
+    node instanceof TestableTreeNodeData
+      ? node.testableMetadata.testable instanceof PackageableElement
+        ? getElementTypeIcon(
+            editorStore,
+            editorStore.graphState.getPackageableElementType(
+              node.testableMetadata.testable,
+            ),
+          )
+        : null
+      : null;
   const resultIcon = getTestableResultIcon(
     getNodeTestableResult(
       node,