diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/NetworkEnvironment.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/NetworkEnvironment.java
index afec6daea46c4..98c61a4d91e17 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/NetworkEnvironment.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/NetworkEnvironment.java
@@ -19,6 +19,8 @@
 package org.apache.flink.runtime.io.network;
 
 import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.metrics.Gauge;
+import org.apache.flink.metrics.MetricGroup;
 import org.apache.flink.runtime.io.network.buffer.BufferPool;
 import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
 import org.apache.flink.runtime.io.network.netty.NettyConfig;
@@ -48,6 +50,10 @@ public class NetworkEnvironment {
 
 	private static final Logger LOG = LoggerFactory.getLogger(NetworkEnvironment.class);
 
+	private static final String METRIC_GROUP_NETWORK = "Network";
+	private static final String METRIC_TOTAL_MEMORY_SEGMENT = "TotalMemorySegments";
+	private static final String METRIC_AVAILABLE_MEMORY_SEGMENT = "AvailableMemorySegments";
+
 	private final Object lock = new Object();
 
 	private final NetworkEnvironmentConfiguration config;
@@ -62,7 +68,10 @@ public class NetworkEnvironment {
 
 	private boolean isShutdown;
 
-	public NetworkEnvironment(NetworkEnvironmentConfiguration config, TaskEventPublisher taskEventPublisher) {
+	public NetworkEnvironment(
+			NetworkEnvironmentConfiguration config,
+			TaskEventPublisher taskEventPublisher,
+			MetricGroup metricGroup) {
 		this.config = checkNotNull(config);
 
 		this.networkBufferPool = new NetworkBufferPool(config.numNetworkBuffers(), config.networkBufferSize());
@@ -78,9 +87,21 @@ public NetworkEnvironment(NetworkEnvironmentConfiguration config, TaskEventPubli
 
 		this.taskEventPublisher = checkNotNull(taskEventPublisher);
 
+		registerNetworkMetrics(metricGroup, networkBufferPool);
+
 		isShutdown = false;
 	}
 
+	private static void registerNetworkMetrics(MetricGroup metricGroup, NetworkBufferPool networkBufferPool) {
+		checkNotNull(metricGroup);
+
+		MetricGroup networkGroup = metricGroup.addGroup(METRIC_GROUP_NETWORK);
+		networkGroup.<Integer, Gauge<Integer>>gauge(METRIC_TOTAL_MEMORY_SEGMENT,
+			networkBufferPool::getTotalNumberOfMemorySegments);
+		networkGroup.<Integer, Gauge<Integer>>gauge(METRIC_AVAILABLE_MEMORY_SEGMENT,
+			networkBufferPool::getNumberOfAvailableMemorySegments);
+	}
+
 	// --------------------------------------------------------------------------------------------
 	//  Properties
 	// --------------------------------------------------------------------------------------------
@@ -93,6 +114,7 @@ public ConnectionManager getConnectionManager() {
 		return connectionManager;
 	}
 
+	@VisibleForTesting
 	public NetworkBufferPool getNetworkBufferPool() {
 		return networkBufferPool;
 	}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/metrics/util/MetricUtils.java b/flink-runtime/src/main/java/org/apache/flink/runtime/metrics/util/MetricUtils.java
index e0adcffad8216..c30ef62470df8 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/metrics/util/MetricUtils.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/metrics/util/MetricUtils.java
@@ -25,15 +25,13 @@
 import org.apache.flink.metrics.Gauge;
 import org.apache.flink.metrics.MetricGroup;
 import org.apache.flink.runtime.clusterframework.BootstrapTools;
-import org.apache.flink.runtime.io.network.NetworkEnvironment;
-import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.clusterframework.types.ResourceID;
 import org.apache.flink.runtime.metrics.MetricNames;
 import org.apache.flink.runtime.metrics.MetricRegistry;
 import org.apache.flink.runtime.metrics.groups.JobManagerMetricGroup;
 import org.apache.flink.runtime.metrics.groups.TaskManagerMetricGroup;
 import org.apache.flink.runtime.rpc.RpcService;
 import org.apache.flink.runtime.rpc.akka.AkkaRpcServiceUtils;
-import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
 import org.apache.flink.util.Preconditions;
 
 import org.slf4j.Logger;
@@ -93,23 +91,19 @@ public static JobManagerMetricGroup instantiateJobManagerMetricGroup(
 
 	public static TaskManagerMetricGroup instantiateTaskManagerMetricGroup(
 			MetricRegistry metricRegistry,
-			TaskManagerLocation taskManagerLocation,
-			NetworkEnvironment network,
+			String hostName,
+			ResourceID resourceID,
 			Optional<Time> systemResourceProbeInterval) {
 		final TaskManagerMetricGroup taskManagerMetricGroup = new TaskManagerMetricGroup(
 			metricRegistry,
-			taskManagerLocation.getHostname(),
-			taskManagerLocation.getResourceID().toString());
+			hostName,
+			resourceID.toString());
 
 		MetricGroup statusGroup = taskManagerMetricGroup.addGroup(METRIC_GROUP_STATUS_NAME);
 
 		// Initialize the TM metrics
 		instantiateStatusMetrics(statusGroup);
 
-		MetricGroup networkGroup = statusGroup
-			.addGroup("Network");
-		instantiateNetworkMetrics(networkGroup, network);
-
 		if (systemResourceProbeInterval.isPresent()) {
 			instantiateSystemMetrics(taskManagerMetricGroup, systemResourceProbeInterval.get());
 		}
@@ -139,15 +133,6 @@ public static RpcService startMetricsRpcService(Configuration configuration, Str
 			new BootstrapTools.FixedThreadPoolExecutorConfiguration(1, 1, threadPriority));
 	}
 
-	private static void instantiateNetworkMetrics(
-		MetricGroup metrics,
-		final NetworkEnvironment network) {
-
-		final NetworkBufferPool networkBufferPool = network.getNetworkBufferPool();
-		metrics.<Integer, Gauge<Integer>>gauge("TotalMemorySegments", networkBufferPool::getTotalNumberOfMemorySegments);
-		metrics.<Integer, Gauge<Integer>>gauge("AvailableMemorySegments", networkBufferPool::getNumberOfAvailableMemorySegments);
-	}
-
 	private static void instantiateClassLoaderMetrics(MetricGroup metrics) {
 		final ClassLoadingMXBean mxBean = ManagementFactory.getClassLoadingMXBean();
 		metrics.<Long, Gauge<Long>>gauge("ClassesLoaded", mxBean::getTotalLoadedClassCount);
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskManagerRunner.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskManagerRunner.java
index 694cbe5fd656a..3df99dd35c090 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskManagerRunner.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskManagerRunner.java
@@ -50,6 +50,7 @@
 import org.apache.flink.runtime.security.SecurityConfiguration;
 import org.apache.flink.runtime.security.SecurityUtils;
 import org.apache.flink.runtime.taskmanager.MemoryLogger;
+import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
 import org.apache.flink.runtime.util.EnvironmentInformation;
 import org.apache.flink.runtime.util.ExecutorThreadFactory;
 import org.apache.flink.runtime.util.Hardware;
@@ -357,19 +358,20 @@ public static TaskExecutor startTaskManager(
 				remoteAddress,
 				localCommunicationOnly);
 
+		TaskManagerMetricGroup taskManagerMetricGroup = MetricUtils.instantiateTaskManagerMetricGroup(
+			metricRegistry,
+			TaskManagerLocation.getHostName(remoteAddress),
+			resourceID,
+			taskManagerServicesConfiguration.getSystemResourceMetricsProbingInterval());
+
 		TaskManagerServices taskManagerServices = TaskManagerServices.fromConfiguration(
 			taskManagerServicesConfiguration,
+			taskManagerMetricGroup,
 			resourceID,
 			rpcService.getExecutor(), // TODO replace this later with some dedicated executor for io.
 			EnvironmentInformation.getSizeOfFreeHeapMemoryWithDefrag(),
 			EnvironmentInformation.getMaxJvmHeapMemory());
 
-		TaskManagerMetricGroup taskManagerMetricGroup = MetricUtils.instantiateTaskManagerMetricGroup(
-			metricRegistry,
-			taskManagerServices.getTaskManagerLocation(),
-			taskManagerServices.getNetworkEnvironment(),
-			taskManagerServicesConfiguration.getSystemResourceMetricsProbingInterval());
-
 		TaskManagerConfiguration taskManagerConfiguration = TaskManagerConfiguration.fromConfiguration(configuration);
 
 		String metricQueryServiceAddress = metricRegistry.getMetricQueryServiceGatewayRpcAddress();
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskManagerServices.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskManagerServices.java
index a67ea3cc6ff2e..a3e906622c83c 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskManagerServices.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskManagerServices.java
@@ -33,6 +33,7 @@
 import org.apache.flink.runtime.io.network.NetworkEnvironment;
 import org.apache.flink.runtime.io.network.TaskEventDispatcher;
 import org.apache.flink.runtime.memory.MemoryManager;
+import org.apache.flink.runtime.metrics.groups.TaskManagerMetricGroup;
 import org.apache.flink.runtime.state.TaskExecutorLocalStateStoresManager;
 import org.apache.flink.runtime.taskexecutor.slot.TaskSlotTable;
 import org.apache.flink.runtime.taskexecutor.slot.TimerService;
@@ -220,9 +221,10 @@ public void shutDown() throws FlinkException {
 	/**
 	 * Creates and returns the task manager services.
 	 *
-	 * @param resourceID resource ID of the task manager
 	 * @param taskManagerServicesConfiguration task manager configuration
-	 * @param taskIOExecutor executor for async IO operations.
+	 * @param taskManagerMetricGroup metric group of the task manager
+	 * @param resourceID resource ID of the task manager
+	 * @param taskIOExecutor executor for async IO operations
 	 * @param freeHeapMemoryWithDefrag an estimate of the size of the free heap memory
 	 * @param maxJvmHeapMemory the maximum JVM heap size
 	 * @return task manager components
@@ -230,6 +232,7 @@ public void shutDown() throws FlinkException {
 	 */
 	public static TaskManagerServices fromConfiguration(
 			TaskManagerServicesConfiguration taskManagerServicesConfiguration,
+			TaskManagerMetricGroup taskManagerMetricGroup,
 			ResourceID resourceID,
 			Executor taskIOExecutor,
 			long freeHeapMemoryWithDefrag,
@@ -241,7 +244,7 @@ public static TaskManagerServices fromConfiguration(
 		final TaskEventDispatcher taskEventDispatcher = new TaskEventDispatcher();
 
 		final NetworkEnvironment network = new NetworkEnvironment(
-			taskManagerServicesConfiguration.getNetworkConfig(), taskEventDispatcher);
+			taskManagerServicesConfiguration.getNetworkConfig(), taskEventDispatcher, taskManagerMetricGroup);
 		network.start();
 
 		final KvStateService kvStateService = KvStateService.fromConfiguration(taskManagerServicesConfiguration);
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/NetworkEnvironmentConfigurationBuilder.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/NetworkEnvironmentConfigurationBuilder.java
deleted file mode 100644
index 4927a8fafcbff..0000000000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/NetworkEnvironmentConfigurationBuilder.java
+++ /dev/null
@@ -1,95 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you 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.
- */
-
-package org.apache.flink.runtime.taskmanager;
-
-import org.apache.flink.runtime.io.network.netty.NettyConfig;
-
-/**
- * Builder for the {@link NetworkEnvironmentConfiguration}.
- */
-public class NetworkEnvironmentConfigurationBuilder {
-
-	private int numNetworkBuffers = 1024;
-
-	private int networkBufferSize = 32 * 1024;
-
-	private int partitionRequestInitialBackoff = 0;
-
-	private int partitionRequestMaxBackoff = 0;
-
-	private int networkBuffersPerChannel = 2;
-
-	private int floatingNetworkBuffersPerGate = 8;
-
-	private boolean isCreditBased = true;
-
-	private NettyConfig nettyConfig;
-
-	public NetworkEnvironmentConfigurationBuilder setNumNetworkBuffers(int numNetworkBuffers) {
-		this.numNetworkBuffers = numNetworkBuffers;
-		return this;
-	}
-
-	public NetworkEnvironmentConfigurationBuilder setNetworkBufferSize(int networkBufferSize) {
-		this.networkBufferSize = networkBufferSize;
-		return this;
-	}
-
-	public NetworkEnvironmentConfigurationBuilder setPartitionRequestInitialBackoff(int partitionRequestInitialBackoff) {
-		this.partitionRequestInitialBackoff = partitionRequestInitialBackoff;
-		return this;
-	}
-
-	public NetworkEnvironmentConfigurationBuilder setPartitionRequestMaxBackoff(int partitionRequestMaxBackoff) {
-		this.partitionRequestMaxBackoff = partitionRequestMaxBackoff;
-		return this;
-	}
-
-	public NetworkEnvironmentConfigurationBuilder setNetworkBuffersPerChannel(int networkBuffersPerChannel) {
-		this.networkBuffersPerChannel = networkBuffersPerChannel;
-		return this;
-	}
-
-	public NetworkEnvironmentConfigurationBuilder setFloatingNetworkBuffersPerGate(int floatingNetworkBuffersPerGate) {
-		this.floatingNetworkBuffersPerGate = floatingNetworkBuffersPerGate;
-		return this;
-	}
-
-	public NetworkEnvironmentConfigurationBuilder setIsCreditBased(boolean isCreditBased) {
-		this.isCreditBased = isCreditBased;
-		return this;
-	}
-
-	public NetworkEnvironmentConfigurationBuilder setNettyConfig(NettyConfig nettyConfig) {
-		this.nettyConfig = nettyConfig;
-		return this;
-	}
-
-	public NetworkEnvironmentConfiguration build() {
-		return new NetworkEnvironmentConfiguration(
-			numNetworkBuffers,
-			networkBufferSize,
-			partitionRequestInitialBackoff,
-			partitionRequestMaxBackoff,
-			networkBuffersPerChannel,
-			floatingNetworkBuffersPerGate,
-			isCreditBased,
-			nettyConfig);
-	}
-}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/TaskManagerLocation.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/TaskManagerLocation.java
index 956a2a2558af9..b3dca5577033f 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/TaskManagerLocation.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/TaskManagerLocation.java
@@ -18,8 +18,6 @@
 
 package org.apache.flink.runtime.taskmanager;
 
-import java.net.InetAddress;
-
 import org.apache.flink.runtime.clusterframework.types.ResourceID;
 import org.apache.flink.util.NetUtils;
 
@@ -28,6 +26,8 @@
 
 import javax.annotation.Nonnull;
 
+import java.net.InetAddress;
+
 import static org.apache.flink.util.Preconditions.checkArgument;
 import static org.apache.flink.util.Preconditions.checkNotNull;
 
@@ -82,30 +82,9 @@ public TaskManagerLocation(ResourceID resourceID, InetAddress inetAddress, int d
 		this.dataPort = dataPort;
 
 		// get FQDN hostname on this TaskManager.
-		String fqdnHostName;
-		try {
-			fqdnHostName = this.inetAddress.getCanonicalHostName();
-		}
-		catch (Throwable t) {
-			LOG.warn("Unable to determine the canonical hostname. Input split assignment (such as " +
-					"for HDFS files) may be non-local when the canonical hostname is missing.");
-			LOG.debug("getCanonicalHostName() Exception:", t);
-			fqdnHostName = this.inetAddress.getHostAddress();
-		}
-		this.fqdnHostName = fqdnHostName;
+		this.fqdnHostName = getFqdnHostName(inetAddress);
 
-		if (this.fqdnHostName.equals(this.inetAddress.getHostAddress())) {
-			// this happens when the name lookup fails, either due to an exception,
-			// or because no hostname can be found for the address
-			// take IP textual representation
-			this.hostName = this.fqdnHostName;
-			LOG.warn("No hostname could be resolved for the IP address {}, using IP address as host name. "
-					+ "Local input split assignment (such as for HDFS files) may be impacted.",
-					this.inetAddress.getHostAddress());
-		}
-		else {
-			this.hostName = NetUtils.getHostnameFromFQDN(this.fqdnHostName);
-		}
+		this.hostName = getHostName(inetAddress);
 
 		this.stringRepresentation = String.format(
 				"%s @ %s (dataPort=%d)", resourceID, fqdnHostName, dataPort);
@@ -185,6 +164,50 @@ public String getHostname() {
 		return hostName;
 	}
 
+	/**
+	 * Gets the fully qualified hostname of the TaskManager based on the network address.
+	 *
+	 * @param inetAddress the network address that the TaskManager binds its sockets to
+	 * @return fully qualified hostname of the TaskManager
+	 */
+	private static String getFqdnHostName(InetAddress inetAddress) {
+		String fqdnHostName;
+		try {
+			fqdnHostName = inetAddress.getCanonicalHostName();
+		} catch (Throwable t) {
+			LOG.warn("Unable to determine the canonical hostname. Input split assignment (such as " +
+				"for HDFS files) may be non-local when the canonical hostname is missing.");
+			LOG.debug("getCanonicalHostName() Exception:", t);
+			fqdnHostName = inetAddress.getHostAddress();
+		}
+
+		return fqdnHostName;
+	}
+
+	/**
+	 * Gets the hostname of the TaskManager based on the network address.
+	 *
+	 * @param inetAddress the network address that the TaskManager binds its sockets to
+	 * @return hostname of the TaskManager
+	 */
+	public static String getHostName(InetAddress inetAddress) {
+		String hostName;
+		String fqdnHostName = getFqdnHostName(inetAddress);
+
+		if (fqdnHostName.equals(inetAddress.getHostAddress())) {
+			// this happens when the name lookup fails, either due to an exception,
+			// or because no hostname can be found for the address
+			// take IP textual representation
+			hostName = fqdnHostName;
+			LOG.warn("No hostname could be resolved for the IP address {}, using IP address as host name. "
+				+ "Local input split assignment (such as for HDFS files) may be impacted.", inetAddress.getHostAddress());
+		} else {
+			hostName = NetUtils.getHostnameFromFQDN(fqdnHostName);
+		}
+
+		return hostName;
+	}
+
 	// --------------------------------------------------------------------------------------------
 	// Utilities
 	// --------------------------------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/NetworkEnvironmentBuilder.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/NetworkEnvironmentBuilder.java
new file mode 100644
index 0000000000000..914433431f774
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/NetworkEnvironmentBuilder.java
@@ -0,0 +1,115 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ */
+
+package org.apache.flink.runtime.io.network;
+
+import org.apache.flink.metrics.MetricGroup;
+import org.apache.flink.runtime.io.network.netty.NettyConfig;
+import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
+import org.apache.flink.runtime.taskmanager.NetworkEnvironmentConfiguration;
+
+/**
+ * Builder for the {@link NetworkEnvironment}.
+ */
+public class NetworkEnvironmentBuilder {
+
+	private int numNetworkBuffers = 1024;
+
+	private int networkBufferSize = 32 * 1024;
+
+	private int partitionRequestInitialBackoff = 0;
+
+	private int partitionRequestMaxBackoff = 0;
+
+	private int networkBuffersPerChannel = 2;
+
+	private int floatingNetworkBuffersPerGate = 8;
+
+	private boolean isCreditBased = true;
+
+	private NettyConfig nettyConfig;
+
+	private TaskEventDispatcher taskEventDispatcher = new TaskEventDispatcher();
+
+	private MetricGroup metricGroup = UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup();
+
+	public NetworkEnvironmentBuilder setNumNetworkBuffers(int numNetworkBuffers) {
+		this.numNetworkBuffers = numNetworkBuffers;
+		return this;
+	}
+
+	public NetworkEnvironmentBuilder setNetworkBufferSize(int networkBufferSize) {
+		this.networkBufferSize = networkBufferSize;
+		return this;
+	}
+
+	public NetworkEnvironmentBuilder setPartitionRequestInitialBackoff(int partitionRequestInitialBackoff) {
+		this.partitionRequestInitialBackoff = partitionRequestInitialBackoff;
+		return this;
+	}
+
+	public NetworkEnvironmentBuilder setPartitionRequestMaxBackoff(int partitionRequestMaxBackoff) {
+		this.partitionRequestMaxBackoff = partitionRequestMaxBackoff;
+		return this;
+	}
+
+	public NetworkEnvironmentBuilder setNetworkBuffersPerChannel(int networkBuffersPerChannel) {
+		this.networkBuffersPerChannel = networkBuffersPerChannel;
+		return this;
+	}
+
+	public NetworkEnvironmentBuilder setFloatingNetworkBuffersPerGate(int floatingNetworkBuffersPerGate) {
+		this.floatingNetworkBuffersPerGate = floatingNetworkBuffersPerGate;
+		return this;
+	}
+
+	public NetworkEnvironmentBuilder setIsCreditBased(boolean isCreditBased) {
+		this.isCreditBased = isCreditBased;
+		return this;
+	}
+
+	public NetworkEnvironmentBuilder setNettyConfig(NettyConfig nettyConfig) {
+		this.nettyConfig = nettyConfig;
+		return this;
+	}
+
+	public NetworkEnvironmentBuilder setTaskEventDispatcher(TaskEventDispatcher taskEventDispatcher) {
+		this.taskEventDispatcher = taskEventDispatcher;
+		return this;
+	}
+
+	public NetworkEnvironmentBuilder setMetricGroup(MetricGroup metricGroup) {
+		this.metricGroup = metricGroup;
+		return this;
+	}
+
+	public NetworkEnvironment build() {
+		return new NetworkEnvironment(
+			new NetworkEnvironmentConfiguration(
+				numNetworkBuffers,
+				networkBufferSize,
+				partitionRequestInitialBackoff,
+				partitionRequestMaxBackoff,
+				networkBuffersPerChannel,
+				floatingNetworkBuffersPerGate,
+				isCreditBased,
+				nettyConfig),
+			taskEventDispatcher,
+			metricGroup);
+	}
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/NetworkEnvironmentTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/NetworkEnvironmentTest.java
index 1d8bae0b8ea96..067b6176d33ee 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/NetworkEnvironmentTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/NetworkEnvironmentTest.java
@@ -30,7 +30,6 @@
 import org.apache.flink.runtime.io.network.partition.consumer.RemoteInputChannel;
 import org.apache.flink.runtime.io.network.partition.consumer.SingleInputGate;
 import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
-import org.apache.flink.runtime.taskmanager.NetworkEnvironmentConfigurationBuilder;
 import org.apache.flink.runtime.taskmanager.NoOpTaskActions;
 import org.apache.flink.runtime.taskmanager.Task;
 
@@ -75,10 +74,9 @@ public static List<Boolean> parameters() {
 	 */
 	@Test
 	public void testRegisterTaskUsesBoundedBuffers() throws Exception {
-		final NetworkEnvironment network = new NetworkEnvironment(new NetworkEnvironmentConfigurationBuilder()
+		final NetworkEnvironment network = new NetworkEnvironmentBuilder()
 			.setIsCreditBased(enableCreditBasedFlowControl)
-			.build(),
-			new TaskEventDispatcher());
+			.build();
 
 		// result partitions
 		ResultPartition rp1 = createResultPartition(ResultPartitionType.PIPELINED, 2);
@@ -181,11 +179,10 @@ public void testRegisterTaskWithInsufficientBuffers() throws Exception {
 	}
 
 	private void testRegisterTaskWithLimitedBuffers(int bufferPoolSize) throws Exception {
-		final NetworkEnvironment network = new NetworkEnvironment(new NetworkEnvironmentConfigurationBuilder()
+		final NetworkEnvironment network = new NetworkEnvironmentBuilder()
 			.setNumNetworkBuffers(bufferPoolSize)
 			.setIsCreditBased(enableCreditBasedFlowControl)
-			.build(),
-			new TaskEventDispatcher());
+			.build();
 
 		final ConnectionManager connManager = createDummyConnectionManager();
 
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionTest.java
index b258ac37748b9..8456e9c7cc74c 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionTest.java
@@ -22,11 +22,10 @@
 import org.apache.flink.runtime.io.disk.iomanager.IOManager;
 import org.apache.flink.runtime.io.disk.iomanager.IOManagerAsync;
 import org.apache.flink.runtime.io.network.NetworkEnvironment;
-import org.apache.flink.runtime.io.network.TaskEventDispatcher;
+import org.apache.flink.runtime.io.network.NetworkEnvironmentBuilder;
 import org.apache.flink.runtime.io.network.buffer.BufferBuilder;
 import org.apache.flink.runtime.io.network.buffer.BufferBuilderTestUtils;
 import org.apache.flink.runtime.io.network.buffer.BufferConsumer;
-import org.apache.flink.runtime.taskmanager.NetworkEnvironmentConfigurationBuilder;
 import org.apache.flink.runtime.taskmanager.NoOpTaskActions;
 import org.apache.flink.runtime.taskmanager.TaskActions;
 
@@ -230,9 +229,8 @@ public void testReleaseMemoryOnPipelinedPartition() throws Exception {
 	 */
 	private void testReleaseMemory(final ResultPartitionType resultPartitionType) throws Exception {
 		final int numAllBuffers = 10;
-		final NetworkEnvironment network = new NetworkEnvironment(new NetworkEnvironmentConfigurationBuilder()
-			.setNumNetworkBuffers(numAllBuffers).build(),
-			new TaskEventDispatcher());
+		final NetworkEnvironment network = new NetworkEnvironmentBuilder()
+			.setNumNetworkBuffers(numAllBuffers).build();
 		final ResultPartitionConsumableNotifier notifier = new NoOpResultPartitionConsumableNotifier();
 		final ResultPartition resultPartition = createPartition(notifier, resultPartitionType, false);
 		try {
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/SingleInputGateTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/SingleInputGateTest.java
index e7aa3cdc3b998..01f60f1d09ea8 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/SingleInputGateTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/SingleInputGateTest.java
@@ -29,6 +29,7 @@
 import org.apache.flink.runtime.io.network.ConnectionManager;
 import org.apache.flink.runtime.io.network.LocalConnectionManager;
 import org.apache.flink.runtime.io.network.NetworkEnvironment;
+import org.apache.flink.runtime.io.network.NetworkEnvironmentBuilder;
 import org.apache.flink.runtime.io.network.TaskEventDispatcher;
 import org.apache.flink.runtime.io.network.buffer.BufferPool;
 import org.apache.flink.runtime.io.network.buffer.FreeingBufferRecycler;
@@ -44,7 +45,6 @@
 import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
 import org.apache.flink.runtime.jobgraph.IntermediateResultPartitionID;
 import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
-import org.apache.flink.runtime.taskmanager.NetworkEnvironmentConfigurationBuilder;
 import org.apache.flink.runtime.taskmanager.NoOpTaskActions;
 
 import org.junit.Test;
@@ -343,12 +343,11 @@ public void testRequestBackoffConfiguration() throws Exception {
 		int initialBackoff = 137;
 		int maxBackoff = 1001;
 
-		final NetworkEnvironment netEnv = new NetworkEnvironment(new NetworkEnvironmentConfigurationBuilder()
+		final NetworkEnvironment netEnv = new NetworkEnvironmentBuilder()
 			.setPartitionRequestInitialBackoff(initialBackoff)
 			.setPartitionRequestMaxBackoff(maxBackoff)
 			.setIsCreditBased(enableCreditBasedFlowControl)
-			.build(),
-			new TaskEventDispatcher());
+			.build();
 
 		SingleInputGate gate = SingleInputGate.create(
 			"TestTask",
@@ -596,10 +595,9 @@ private void addRemoteInputChannel(
 	}
 
 	private NetworkEnvironment createNetworkEnvironment() {
-		return new NetworkEnvironment(new NetworkEnvironmentConfigurationBuilder()
+		return new NetworkEnvironmentBuilder()
 			.setIsCreditBased(enableCreditBasedFlowControl)
-			.build(),
-			new TaskEventDispatcher());
+			.build();
 	}
 
 	static void verifyBufferOrEvent(
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManagerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManagerTest.java
index e99522360a7d8..c4c74310ccfa5 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManagerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManagerTest.java
@@ -25,6 +25,7 @@
 import org.apache.flink.runtime.clusterframework.types.ResourceID;
 import org.apache.flink.runtime.concurrent.Executors;
 import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
 import org.apache.flink.runtime.taskexecutor.TaskManagerServices;
 import org.apache.flink.runtime.taskexecutor.TaskManagerServicesConfiguration;
 import org.apache.flink.util.FileUtils;
@@ -36,6 +37,7 @@
 import org.junit.rules.TemporaryFolder;
 
 import java.io.File;
+import java.io.IOException;
 import java.net.InetAddress;
 
 public class TaskExecutorLocalStateStoresManagerTest extends TestLogger {
@@ -64,17 +66,7 @@ public void testCreationFromConfig() throws Exception {
 		// test configuration of the local state mode
 		config.setBoolean(CheckpointingOptions.LOCAL_RECOVERY, true);
 
-		final ResourceID tmResourceID = ResourceID.generate();
-
-		TaskManagerServicesConfiguration taskManagerServicesConfiguration =
-			TaskManagerServicesConfiguration.fromConfiguration(config, MEM_SIZE_PARAM, InetAddress.getLocalHost(), true);
-
-		TaskManagerServices taskManagerServices = TaskManagerServices.fromConfiguration(
-			taskManagerServicesConfiguration,
-			tmResourceID,
-			Executors.directExecutor(),
-			MEM_SIZE_PARAM,
-			MEM_SIZE_PARAM);
+		TaskManagerServices taskManagerServices = createTaskManagerServices(createTaskManagerServiceConfiguration(config));
 
 		TaskExecutorLocalStateStoresManager taskStateManager = taskManagerServices.getTaskManagerStateStore();
 
@@ -105,17 +97,9 @@ public void testCreationFromConfigDefault() throws Exception {
 
 		final Configuration config = new Configuration();
 
-		final ResourceID tmResourceID = ResourceID.generate();
-
-		TaskManagerServicesConfiguration taskManagerServicesConfiguration =
-			TaskManagerServicesConfiguration.fromConfiguration(config, MEM_SIZE_PARAM, InetAddress.getLocalHost(), true);
+		TaskManagerServicesConfiguration taskManagerServicesConfiguration = createTaskManagerServiceConfiguration(config);
 
-		TaskManagerServices taskManagerServices = TaskManagerServices.fromConfiguration(
-			taskManagerServicesConfiguration,
-			tmResourceID,
-			Executors.directExecutor(),
-			MEM_SIZE_PARAM,
-			MEM_SIZE_PARAM);
+		TaskManagerServices taskManagerServices = createTaskManagerServices(taskManagerServicesConfiguration);
 
 		TaskExecutorLocalStateStoresManager taskStateManager = taskManagerServices.getTaskManagerStateStore();
 
@@ -217,4 +201,21 @@ private void checkRootDirsClean(File[] rootDirs) {
 			}
 		}
 	}
+
+	private TaskManagerServicesConfiguration createTaskManagerServiceConfiguration(
+			Configuration config) throws IOException {
+		return TaskManagerServicesConfiguration.fromConfiguration(
+			config, MEM_SIZE_PARAM, InetAddress.getLocalHost(), true);
+	}
+
+	private TaskManagerServices createTaskManagerServices(
+			TaskManagerServicesConfiguration config) throws Exception {
+		return TaskManagerServices.fromConfiguration(
+			config,
+			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
+			ResourceID.generate(),
+			Executors.directExecutor(),
+			MEM_SIZE_PARAM,
+			MEM_SIZE_PARAM);
+	}
 }
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorTest.java
index 2076c20a0ada0..b47ad53284547 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorTest.java
@@ -57,7 +57,7 @@
 import org.apache.flink.runtime.io.disk.iomanager.IOManager;
 import org.apache.flink.runtime.io.disk.iomanager.IOManagerAsync;
 import org.apache.flink.runtime.io.network.NetworkEnvironment;
-import org.apache.flink.runtime.io.network.TaskEventDispatcher;
+import org.apache.flink.runtime.io.network.NetworkEnvironmentBuilder;
 import org.apache.flink.runtime.io.network.netty.PartitionProducerStateChecker;
 import org.apache.flink.runtime.io.network.partition.NoOpResultPartitionConsumableNotifier;
 import org.apache.flink.runtime.jobgraph.JobVertexID;
@@ -90,7 +90,6 @@
 import org.apache.flink.runtime.taskexecutor.slot.TimerService;
 import org.apache.flink.runtime.taskmanager.CheckpointResponder;
 import org.apache.flink.runtime.taskmanager.LocalTaskManagerLocation;
-import org.apache.flink.runtime.taskmanager.NetworkEnvironmentConfigurationBuilder;
 import org.apache.flink.runtime.taskmanager.NoOpTaskManagerActions;
 import org.apache.flink.runtime.taskmanager.Task;
 import org.apache.flink.runtime.taskmanager.TaskExecutionState;
@@ -263,8 +262,7 @@ public void testShouldShutDownTaskManagerServicesInPostStop() throws Exception {
 			MemoryType.HEAP,
 			false);
 
-		final NetworkEnvironment networkEnvironment = new NetworkEnvironment(
-			new NetworkEnvironmentConfigurationBuilder().build(), new TaskEventDispatcher());
+		final NetworkEnvironment networkEnvironment = new NetworkEnvironmentBuilder().build();
 		networkEnvironment.start();
 
 		final KvStateService kvStateService = new KvStateService(new KvStateRegistry(), null, null);
@@ -281,20 +279,7 @@ public void testShouldShutDownTaskManagerServicesInPostStop() throws Exception {
 			.setTaskStateManager(localStateStoresManager)
 			.build();
 
-		final long heartbeatInterval = 1000L;
-		final long heartbeatTimeout = 1000L;
-		final HeartbeatServices heartbeatServices = new HeartbeatServices(heartbeatInterval, heartbeatTimeout);
-
-		final TaskExecutor taskManager = new TaskExecutor(
-			rpc,
-			taskManagerConfiguration,
-			haServices,
-			taskManagerServices,
-			heartbeatServices,
-			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
-			null,
-			dummyBlobCacheService,
-			testingFatalErrorHandler);
+		final TaskExecutor taskManager = createTaskExecutor(taskManagerServices);
 
 		try {
 			taskManager.start();
@@ -349,16 +334,7 @@ public void testHeartbeatTimeoutWithJobManager() throws Exception {
 			.setTaskStateManager(localStateStoresManager)
 			.build();
 
-		final TestingTaskExecutor taskManager = new TestingTaskExecutor(
-			rpc,
-			taskManagerConfiguration,
-			haServices,
-			taskManagerServices,
-			heartbeatServices,
-			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
-			null,
-			dummyBlobCacheService,
-			testingFatalErrorHandler);
+		final TestingTaskExecutor taskManager = createTestingTaskExecutor(taskManagerServices, heartbeatServices);
 
 		try {
 			taskManager.start();
@@ -439,16 +415,7 @@ public void testHeartbeatTimeoutWithResourceManager() throws Exception {
 			.setTaskStateManager(localStateStoresManager)
 			.build();
 
-		final TaskExecutor taskManager = new TaskExecutor(
-			rpc,
-			taskManagerConfiguration,
-			haServices,
-			taskManagerServices,
-			heartbeatServices,
-			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
-			null,
-			dummyBlobCacheService,
-			testingFatalErrorHandler);
+		final TaskExecutor taskManager = createTaskExecutor(taskManagerServices, heartbeatServices);
 
 		try {
 			taskManager.start();
@@ -527,16 +494,7 @@ public void testHeartbeatSlotReporting() throws Exception {
 			.setTaskStateManager(localStateStoresManager)
 			.build();
 
-		final TaskExecutor taskManager = new TaskExecutor(
-			rpc,
-			taskManagerConfiguration,
-			haServices,
-			taskManagerServices,
-			HEARTBEAT_SERVICES,
-			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
-			null,
-			dummyBlobCacheService,
-			testingFatalErrorHandler);
+		final TaskExecutor taskManager = createTaskExecutor(taskManagerServices);
 
 		try {
 			taskManager.start();
@@ -636,16 +594,7 @@ public void testTriggerRegistrationOnLeaderChange() throws Exception {
 			.setTaskStateManager(localStateStoresManager)
 			.build();
 
-		TaskExecutor taskManager = new TaskExecutor(
-			rpc,
-			taskManagerConfiguration,
-			haServices,
-			taskManagerServices,
-			HEARTBEAT_SERVICES,
-			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
-			null,
-			dummyBlobCacheService,
-			testingFatalErrorHandler);
+		TaskExecutor taskManager = createTaskExecutor(taskManagerServices);
 
 		try {
 			taskManager.start();
@@ -761,16 +710,7 @@ public void updateTaskExecutionState(TaskExecutionState taskExecutionState) {
 			.setTaskStateManager(localStateStoresManager)
 			.build();
 
-		TaskExecutor taskManager = new TaskExecutor(
-			rpc,
-			taskManagerConfiguration,
-			haServices,
-			taskManagerServices,
-			HEARTBEAT_SERVICES,
-			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
-			null,
-			dummyBlobCacheService,
-			testingFatalErrorHandler);
+		TaskExecutor taskManager = createTaskExecutor(taskManagerServices);
 
 		try {
 			taskManager.start();
@@ -848,16 +788,7 @@ public void testJobLeaderDetection() throws Exception {
 			.setTaskStateManager(localStateStoresManager)
 			.build();
 
-		TaskExecutor taskManager = new TaskExecutor(
-			rpc,
-			taskManagerConfiguration,
-			haServices,
-			taskManagerServices,
-			HEARTBEAT_SERVICES,
-			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
-			null,
-			dummyBlobCacheService,
-			testingFatalErrorHandler);
+		TaskExecutor taskManager = createTaskExecutor(taskManagerServices);
 
 		try {
 			taskManager.start();
@@ -958,16 +889,7 @@ public void testSlotAcceptance() throws Exception {
 			.setTaskStateManager(localStateStoresManager)
 			.build();
 
-		TaskExecutor taskManager = new TaskExecutor(
-			rpc,
-			taskManagerConfiguration,
-			haServices,
-			taskManagerServices,
-			HEARTBEAT_SERVICES,
-			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
-			null,
-			dummyBlobCacheService,
-			testingFatalErrorHandler);
+		TaskExecutor taskManager = createTaskExecutor(taskManagerServices);
 
 		try {
 			taskManager.start();
@@ -1049,16 +971,7 @@ public void testSubmitTaskBeforeAcceptSlot() throws Exception {
 			.setTaskStateManager(localStateStoresManager)
 			.build();
 
-		final TestingTaskExecutor taskManager = new TestingTaskExecutor(
-			rpc,
-			taskManagerConfiguration,
-			haServices,
-			taskManagerServices,
-			HEARTBEAT_SERVICES,
-			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
-			null,
-			dummyBlobCacheService,
-			testingFatalErrorHandler);
+		final TestingTaskExecutor taskManager = createTestingTaskExecutor(taskManagerServices);
 
 		try {
 			taskManager.start();
@@ -1137,7 +1050,6 @@ public void testSubmitTaskBeforeAcceptSlot() throws Exception {
 	public void testFilterOutDuplicateJobMasterRegistrations() throws Exception {
 		final long verificationTimeout = 500L;
 		final JobLeaderService jobLeaderService = mock(JobLeaderService.class);
-		final HeartbeatServices heartbeatServicesMock = mock(HeartbeatServices.class, Mockito.RETURNS_MOCKS);
 
 		final JobMasterGateway jobMasterGateway = mock(JobMasterGateway.class);
 		when(jobMasterGateway.getHostname()).thenReturn("localhost");
@@ -1153,16 +1065,7 @@ public void testFilterOutDuplicateJobMasterRegistrations() throws Exception {
 			.setTaskStateManager(localStateStoresManager)
 			.build();
 
-		final TestingTaskExecutor taskExecutor = new TestingTaskExecutor(
-			rpc,
-			taskManagerConfiguration,
-			haServices,
-			taskManagerServices,
-			heartbeatServicesMock,
-			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
-			null,
-			dummyBlobCacheService,
-			testingFatalErrorHandler);
+		final TestingTaskExecutor taskExecutor = createTestingTaskExecutor(taskManagerServices);
 
 		try {
 			taskExecutor.start();
@@ -1223,16 +1126,7 @@ public void testRMHeartbeatStopWhenLeadershipRevoked() throws Exception {
 			.setTaskStateManager(localStateStoresManager)
 			.build();
 
-		final TaskExecutor taskExecutor = new TaskExecutor(
-			rpc,
-			taskManagerConfiguration,
-			haServices,
-			taskManagerServices,
-			heartbeatServices,
-			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
-			null,
-			dummyBlobCacheService,
-			testingFatalErrorHandler);
+		final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices, heartbeatServices);
 
 		try {
 			taskExecutor.start();
@@ -1277,16 +1171,7 @@ public void testRemoveJobFromJobLeaderService() throws Exception {
 			.setTaskStateManager(localStateStoresManager)
 			.build();
 
-		final TestingTaskExecutor taskExecutor = new TestingTaskExecutor(
-			rpc,
-			taskManagerConfiguration,
-			haServices,
-			taskManagerServices,
-			HEARTBEAT_SERVICES,
-			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
-			null,
-			dummyBlobCacheService,
-			testingFatalErrorHandler);
+		final TestingTaskExecutor taskExecutor = createTestingTaskExecutor(taskManagerServices);
 
 		try {
 			final TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
@@ -1369,18 +1254,8 @@ public void testMaximumRegistrationDurationAfterConnectionLoss() throws Exceptio
 		configuration.setString(TaskManagerOptions.REGISTRATION_TIMEOUT, "100 ms");
 		final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
 
-		final long heartbeatInterval = 10L;
 		final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setTaskSlotTable(taskSlotTable).build();
-		final TaskExecutor taskExecutor = new TaskExecutor(
-			rpc,
-			TaskManagerConfiguration.fromConfiguration(configuration),
-			haServices,
-			taskManagerServices,
-			new HeartbeatServices(heartbeatInterval, 10L),
-			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
-			null,
-			dummyBlobCacheService,
-			testingFatalErrorHandler);
+		final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices, new HeartbeatServices(10L, 10L));
 
 		taskExecutor.start();
 
@@ -1472,22 +1347,12 @@ public void testIgnoringSlotRequestsIfNotRegistered() throws Exception {
 	 */
 	@Test
 	public void testReconnectionAttemptIfExplicitlyDisconnected() throws Exception {
-		final long heartbeatInterval = 1000L;
 		final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
 		final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
-		final TaskExecutor taskExecutor = new TaskExecutor(
-			rpc,
-			TaskManagerConfiguration.fromConfiguration(configuration),
-			haServices,
-			new TaskManagerServicesBuilder()
-				.setTaskSlotTable(taskSlotTable)
-				.setTaskManagerLocation(taskManagerLocation)
-				.build(),
-			new HeartbeatServices(heartbeatInterval, 1000L),
-			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
-			null,
-			dummyBlobCacheService,
-			testingFatalErrorHandler);
+		final TaskExecutor taskExecutor = createTaskExecutor(new TaskManagerServicesBuilder()
+			.setTaskSlotTable(taskSlotTable)
+			.setTaskManagerLocation(taskManagerLocation)
+			.build());
 
 		taskExecutor.start();
 
@@ -1803,12 +1668,33 @@ private TaskExecutorLocalStateStoresManager createTaskExecutorLocalStateStoresMa
 
 	@Nonnull
 	private TaskExecutor createTaskExecutor(TaskManagerServices taskManagerServices) {
+		return createTaskExecutor(taskManagerServices, HEARTBEAT_SERVICES);
+	}
+
+	private TaskExecutor createTaskExecutor(TaskManagerServices taskManagerServices, HeartbeatServices heartbeatServices) {
 		return new TaskExecutor(
 			rpc,
 			TaskManagerConfiguration.fromConfiguration(configuration),
 			haServices,
 			taskManagerServices,
-			HEARTBEAT_SERVICES,
+			heartbeatServices,
+			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
+			null,
+			dummyBlobCacheService,
+			testingFatalErrorHandler);
+	}
+
+	private TestingTaskExecutor createTestingTaskExecutor(TaskManagerServices taskManagerServices) {
+		return createTestingTaskExecutor(taskManagerServices, HEARTBEAT_SERVICES);
+	}
+
+	private TestingTaskExecutor createTestingTaskExecutor(TaskManagerServices taskManagerServices, HeartbeatServices heartbeatServices) {
+		return new TestingTaskExecutor(
+			rpc,
+			TaskManagerConfiguration.fromConfiguration(configuration),
+			haServices,
+			taskManagerServices,
+			heartbeatServices,
 			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
 			null,
 			dummyBlobCacheService,
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskSubmissionTestEnvironment.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskSubmissionTestEnvironment.java
index 689d212c3a37f..0f298805aec7a 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskSubmissionTestEnvironment.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskSubmissionTestEnvironment.java
@@ -35,7 +35,7 @@
 import org.apache.flink.runtime.heartbeat.HeartbeatServices;
 import org.apache.flink.runtime.highavailability.TestingHighAvailabilityServices;
 import org.apache.flink.runtime.io.network.NetworkEnvironment;
-import org.apache.flink.runtime.io.network.TaskEventDispatcher;
+import org.apache.flink.runtime.io.network.NetworkEnvironmentBuilder;
 import org.apache.flink.runtime.io.network.netty.NettyConfig;
 import org.apache.flink.runtime.io.network.netty.PartitionProducerStateChecker;
 import org.apache.flink.runtime.jobmaster.JobMasterGateway;
@@ -53,7 +53,6 @@
 import org.apache.flink.runtime.taskexecutor.slot.TimerService;
 import org.apache.flink.runtime.taskmanager.CheckpointResponder;
 import org.apache.flink.runtime.taskmanager.NetworkEnvironmentConfiguration;
-import org.apache.flink.runtime.taskmanager.NetworkEnvironmentConfigurationBuilder;
 import org.apache.flink.runtime.taskmanager.NoOpTaskManagerActions;
 import org.apache.flink.runtime.taskmanager.Task;
 import org.apache.flink.runtime.taskmanager.TaskManagerActions;
@@ -246,12 +245,11 @@ private static NetworkEnvironment createNetworkEnvironment(
 			final NettyConfig nettyConfig = new NettyConfig(socketAddress.getAddress(), socketAddress.getPort(),
 				NetworkEnvironmentConfiguration.getPageSize(configuration), ConfigurationParserUtils.getSlot(configuration), configuration);
 
-			networkEnvironment =  new NetworkEnvironment(new NetworkEnvironmentConfigurationBuilder()
+			networkEnvironment =  new NetworkEnvironmentBuilder()
 				.setPartitionRequestInitialBackoff(configuration.getInteger(TaskManagerOptions.NETWORK_REQUEST_BACKOFF_INITIAL))
 				.setPartitionRequestMaxBackoff(configuration.getInteger(TaskManagerOptions.NETWORK_REQUEST_BACKOFF_MAX))
 				.setNettyConfig(localCommunication ? null : nettyConfig)
-				.build(),
-				new TaskEventDispatcher());
+				.build();
 			networkEnvironment.start();
 		}
 
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/io/benchmark/StreamNetworkBenchmarkEnvironment.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/io/benchmark/StreamNetworkBenchmarkEnvironment.java
index 87dfa138928d8..0edf2b09feca9 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/io/benchmark/StreamNetworkBenchmarkEnvironment.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/io/benchmark/StreamNetworkBenchmarkEnvironment.java
@@ -30,6 +30,7 @@
 import org.apache.flink.runtime.io.disk.iomanager.IOManagerAsync;
 import org.apache.flink.runtime.io.network.ConnectionID;
 import org.apache.flink.runtime.io.network.NetworkEnvironment;
+import org.apache.flink.runtime.io.network.NetworkEnvironmentBuilder;
 import org.apache.flink.runtime.io.network.TaskEventDispatcher;
 import org.apache.flink.runtime.io.network.api.writer.RecordWriter;
 import org.apache.flink.runtime.io.network.api.writer.RecordWriterBuilder;
@@ -45,7 +46,6 @@
 import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
 import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
 import org.apache.flink.runtime.taskmanager.NetworkEnvironmentConfiguration;
-import org.apache.flink.runtime.taskmanager.NetworkEnvironmentConfigurationBuilder;
 import org.apache.flink.runtime.taskmanager.NoOpTaskActions;
 import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
 import org.apache.flink.runtime.util.ConfigurationParserUtils;
@@ -198,12 +198,10 @@ private NetworkEnvironment createNettyNetworkEnvironment(
 			// please note that the number of slots directly influences the number of netty threads!
 			ConfigurationParserUtils.getSlot(config),
 			config);
-		final NetworkEnvironmentConfiguration configuration = new NetworkEnvironmentConfigurationBuilder()
+		return new NetworkEnvironmentBuilder()
 			.setNumNetworkBuffers(bufferPoolSize)
 			.setNettyConfig(nettyConfig)
 			.build();
-
-		return new NetworkEnvironment(configuration, new TaskEventDispatcher());
 	}
 
 	protected ResultPartitionWriter createResultPartition(