diff --git a/single-line-diagram/single-line-diagram-core/src/test/java/com/powsybl/sld/iidm/TestUnifyVoltageLevelColors.java b/single-line-diagram/single-line-diagram-core/src/test/java/com/powsybl/sld/iidm/TestUnifyVoltageLevelColors.java
index 896e46954..576fd0dd5 100644
--- a/single-line-diagram/single-line-diagram-core/src/test/java/com/powsybl/sld/iidm/TestUnifyVoltageLevelColors.java
+++ b/single-line-diagram/single-line-diagram-core/src/test/java/com/powsybl/sld/iidm/TestUnifyVoltageLevelColors.java
@@ -28,16 +28,17 @@ class TestUnifyVoltageLevelColors extends AbstractTestCaseIidm {
 
     @Override
     public void setUp() throws IOException {
+        network = Networks.createNetworkWithTieLineInSubstation();
+        graphBuilder = new NetworkGraphBuilder(network);
+        substation = network.getSubstation("S1");
     }
 
     @Test
     void testSubstationRegularVoltageLevelColors() {
         svgParameters.setUnifyVoltageLevelColors(false);
 
-        Network network = Networks.createNetworkWithTieLineInSubstation();
-        graphBuilder = new NetworkGraphBuilder(network);
         // build graph
-        SubstationGraph g = graphBuilder.buildSubstationGraph("S1");
+        SubstationGraph g = graphBuilder.buildSubstationGraph(substation.getId());
 
         // Run layout
         new VerticalSubstationLayoutFactory().create(g, new SmartVoltageLevelLayoutFactory(network)).run(layoutParameters);
@@ -50,10 +51,8 @@ void testSubstationRegularVoltageLevelColors() {
     void testSubstationUnifiedVoltageLevelColors() {
         svgParameters.setUnifyVoltageLevelColors(true);
 
-        Network network = Networks.createNetworkWithTieLineInSubstation();
-        graphBuilder = new NetworkGraphBuilder(network);
         // build graph
-        SubstationGraph g = graphBuilder.buildSubstationGraph("S1");
+        SubstationGraph g = graphBuilder.buildSubstationGraph(substation.getId());
 
         // Run layout
         new VerticalSubstationLayoutFactory().create(g, new SmartVoltageLevelLayoutFactory(network)).run(layoutParameters);