diff --git a/packages/serverless-workflow-diagram-editor/kie-wb-common-stunner/kie-wb-common-stunner-core/kie-wb-common-stunner-commons/kie-wb-common-stunner-client-common/src/test/java/org/kie/workbench/common/stunner/core/client/shape/impl/ShapeStateAttributesFactoryTest.java b/packages/serverless-workflow-diagram-editor/kie-wb-common-stunner/kie-wb-common-stunner-core/kie-wb-common-stunner-commons/kie-wb-common-stunner-client-common/src/test/java/org/kie/workbench/common/stunner/core/client/shape/impl/ShapeStateAttributesFactoryTest.java
index 800a6932820..576d2adadce 100644
--- a/packages/serverless-workflow-diagram-editor/kie-wb-common-stunner/kie-wb-common-stunner-core/kie-wb-common-stunner-commons/kie-wb-common-stunner-client-common/src/test/java/org/kie/workbench/common/stunner/core/client/shape/impl/ShapeStateAttributesFactoryTest.java
+++ b/packages/serverless-workflow-diagram-editor/kie-wb-common-stunner/kie-wb-common-stunner-core/kie-wb-common-stunner-commons/kie-wb-common-stunner-client-common/src/test/java/org/kie/workbench/common/stunner/core/client/shape/impl/ShapeStateAttributesFactoryTest.java
@@ -20,61 +20,52 @@
 
 package org.kie.workbench.common.stunner.core.client.shape.impl;
 
-import java.util.function.Function;
-
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.kie.workbench.common.stunner.core.client.shape.ShapeState;
-import org.kie.workbench.common.stunner.core.client.theme.StunnerColorTheme;
 import org.kie.workbench.common.stunner.core.client.theme.StunnerTheme;
 import org.mockito.junit.MockitoJUnitRunner;
 
 import static org.junit.Assert.assertEquals;
-import static org.kie.workbench.common.stunner.core.client.shape.ShapeState.HIGHLIGHT;
-import static org.kie.workbench.common.stunner.core.client.shape.ShapeState.INVALID;
-import static org.kie.workbench.common.stunner.core.client.shape.ShapeState.NONE;
-import static org.kie.workbench.common.stunner.core.client.shape.ShapeState.SELECTED;
-import static org.kie.workbench.common.stunner.core.client.shape.impl.ShapeStateAttributeHandler.ShapeStateAttribute.FILL_ALPHA;
-import static org.kie.workbench.common.stunner.core.client.shape.impl.ShapeStateAttributeHandler.ShapeStateAttribute.FILL_COLOR;
-import static org.kie.workbench.common.stunner.core.client.shape.impl.ShapeStateAttributeHandler.ShapeStateAttribute.STROKE_ALPHA;
-import static org.kie.workbench.common.stunner.core.client.shape.impl.ShapeStateAttributeHandler.ShapeStateAttribute.STROKE_COLOR;
 
 @RunWith(MockitoJUnitRunner.class)
 public class ShapeStateAttributesFactoryTest {
 
-    private StunnerColorTheme theme = StunnerTheme.getTheme();
-    private Function<ShapeState, ShapeStateAttributeHandler.ShapeStateAttributes> stateAttributes = ShapeStateAttributesFactory::buildStateAttributes;
-
     @Test
     public void testBuildStrokeAttributes() {
-        ShapeStateAttributeHandler.ShapeStateAttributes attributes = stateAttributes.apply(NONE);
+        ShapeStateAttributeHandler.ShapeStateAttributes attributes = ShapeStateAttributesFactory.buildStateAttributes(ShapeState.NONE);
         attributes.getValues().values().forEach(
                 state -> assertEquals(null, state)
         );
 
-        assertStrokeAttributes(stateAttributes.apply(SELECTED), theme.getShapeStrokeColorSelected());
-        assertStrokeAttributes(stateAttributes.apply(HIGHLIGHT), theme.getShapeStrokeColorHighlight());
-        assertStrokeAttributes(stateAttributes.apply(INVALID), theme.getShapeStrokeColorInvalid());
+        assertStrokeAttributes(ShapeStateAttributesFactory.buildStateAttributes(ShapeState.SELECTED),
+                               StunnerTheme.getTheme().getShapeStrokeColorSelected());
+        assertStrokeAttributes(ShapeStateAttributesFactory.buildStateAttributes(ShapeState.HIGHLIGHT),
+                               StunnerTheme.getTheme().getShapeStrokeColorHighlight());
+        assertStrokeAttributes(ShapeStateAttributesFactory.buildStateAttributes(ShapeState.INVALID),
+                               StunnerTheme.getTheme().getShapeStrokeColorInvalid());
     }
 
     @Test
     public void testBuildFillAttributes() {
-        ShapeStateAttributeHandler.ShapeStateAttributes attributes = stateAttributes.apply(NONE);
+        ShapeStateAttributeHandler.ShapeStateAttributes attributes = ShapeStateAttributesFactory.buildStateAttributes(ShapeState.NONE);
         attributes.getValues().values().forEach(
                 state -> assertEquals(null, state)
         );
 
-        assertFillAttributes(stateAttributes.apply(SELECTED), theme.getShapeFillColorSelected());
-        assertFillAttributes(stateAttributes.apply(INVALID), theme.getShapeFillColorInvalid());
+        assertFillAttributes(ShapeStateAttributesFactory.buildStateAttributes(ShapeState.SELECTED),
+                             StunnerTheme.getTheme().getShapeFillColorSelected());
+        assertFillAttributes(ShapeStateAttributesFactory.buildStateAttributes(ShapeState.INVALID),
+                             StunnerTheme.getTheme().getShapeFillColorInvalid());
     }
 
     private void assertStrokeAttributes(ShapeStateAttributeHandler.ShapeStateAttributes attributes, String color) {
-        assertEquals(1d, attributes.getValues().get(STROKE_ALPHA));
-        assertEquals(color, attributes.getValues().get(STROKE_COLOR));
+        assertEquals(1d, attributes.getValues().get(ShapeStateAttributeHandler.ShapeStateAttribute.STROKE_ALPHA));
+        assertEquals(color, attributes.getValues().get(ShapeStateAttributeHandler.ShapeStateAttribute.STROKE_COLOR));
     }
 
     private void assertFillAttributes(ShapeStateAttributeHandler.ShapeStateAttributes attributes, String color) {
-        assertEquals(1d, attributes.getValues().get(FILL_ALPHA));
-        assertEquals(color, attributes.getValues().get(FILL_COLOR));
+        assertEquals(1d, attributes.getValues().get(ShapeStateAttributeHandler.ShapeStateAttribute.FILL_ALPHA));
+        assertEquals(color, attributes.getValues().get(ShapeStateAttributeHandler.ShapeStateAttribute.FILL_COLOR));
     }
 }