diff --git a/libs/ui/dataviz/src/lib/figure/figure.component.html b/libs/ui/dataviz/src/lib/figure/figure.component.html
index 285d3f914b..5e4f8511ff 100644
--- a/libs/ui/dataviz/src/lib/figure/figure.component.html
+++ b/libs/ui/dataviz/src/lib/figure/figure.component.html
@@ -14,6 +14,7 @@
       textClass
     }} text-[1.875em] rounded-full mr-[0.55em] flex justify-center items-center w-[2.2em] h-[2.2em] shrink-0"
     style="width: 2.2em; height: 2.2em"
+    data-test="icon-container"
   >
     <ng-icon class="text-[0.66em]" [name]="icon"> </ng-icon>
   </div>
diff --git a/libs/ui/dataviz/src/lib/figure/figure.component.spec.ts b/libs/ui/dataviz/src/lib/figure/figure.component.spec.ts
index adaef407eb..86f0bc9034 100644
--- a/libs/ui/dataviz/src/lib/figure/figure.component.spec.ts
+++ b/libs/ui/dataviz/src/lib/figure/figure.component.spec.ts
@@ -54,12 +54,14 @@ describe('FigureComponent', () => {
     expect(icon['name']).toContain('group')
   })
   it('icon is primary color', () => {
-    expect(compiled.querySelector('ng-icon')?.className).toContain(
-      'text-primary'
-    )
+    expect(
+      compiled.querySelector('[data-test=icon-container]')?.className
+    ).toContain('text-primary')
   })
   it('icon background is primary color', () => {
-    expect(compiled.querySelector('ng-icon')?.className).toContain('bg-primary')
+    expect(
+      compiled.querySelector('[data-test=icon-container]')?.className
+    ).toContain('bg-primary')
   })
   it('label is black text color', () => {
     expect(compiled.querySelector('.figure-block')?.className).toContain(
@@ -95,14 +97,14 @@ describe('FigureComponent', () => {
       fixture.detectChanges()
     })
     it('icon is secondary color', () => {
-      expect(compiled.querySelector('ng-icon')?.className).toContain(
-        'text-secondary'
-      )
+      expect(
+        compiled.querySelector('[data-test=icon-container]')?.className
+      ).toContain('text-secondary')
     })
     it('icon background is secondary color', () => {
-      expect(compiled.querySelector('ng-icon')?.className).toContain(
-        'bg-secondary'
-      )
+      expect(
+        compiled.querySelector('[data-test=icon-container]')?.className
+      ).toContain('bg-secondary')
     })
   })
 })