diff --git a/components/esp-tls/test/test_esp_tls.c b/components/esp-tls/test/test_esp_tls.c
index 3635baaee1..2d22e31e49 100644
--- a/components/esp-tls/test/test_esp_tls.c
+++ b/components/esp-tls/test/test_esp_tls.c
@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-#include "test_utils.h"
+#include "memory_checks.h"
 #include "esp_tls.h"
 #include "unity.h"
 #include "esp_err.h"
@@ -75,7 +75,7 @@ static void test_leak_setup(const char *file, long line)
     const uint8_t input_buffer[64];
     uint8_t output_buffer[64];
     esp_sha(SHA2_512, input_buffer, sizeof(input_buffer), output_buffer);
-    unity_reset_leak_checks();
+    test_utils_record_free_mem();
 }
 
 
diff --git a/components/esp_websocket_client/test/test_websocket_client.c b/components/esp_websocket_client/test/test_websocket_client.c
index 5d96b49831..1c9e895a38 100644
--- a/components/esp_websocket_client/test/test_websocket_client.c
+++ b/components/esp_websocket_client/test/test_websocket_client.c
@@ -1,28 +1,26 @@
-// Copyright 2021 Espressif Systems (Shanghai) PTE LTD
-//
-// 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.
+/*
+ * SPDX-FileCopyrightText: 2021 Espressif Systems (Shanghai) CO LTD
+ *
+ * SPDX-License-Identifier: Unlicense OR CC0-1.0
+ *
+ * This test code is in the Public Domain (or CC0 licensed, at your option.)
+ *
+ * Unless required by applicable law or agreed to in writing, this
+ * software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+ * CONDITIONS OF ANY KIND, either express or implied.
+ */
 
 #include <stdlib.h>
 #include <stdbool.h>
 #include <esp_websocket_client.h>
 
 #include "unity.h"
-#include "test_utils.h"
+#include "memory_checks.h"
 
 static void test_leak_setup(const char * file, long line)
 {
     printf("%s:%ld\n", file, line);
-    unity_reset_leak_checks();
+    test_utils_record_free_mem();
 }
 
 TEST_CASE("websocket init and deinit", "[websocket][leaks=0]")
diff --git a/components/esp_wifi/test/test_wifi.c b/components/esp_wifi/test/test_wifi.c
index 56e5930bdf..b3125eb8a7 100644
--- a/components/esp_wifi/test/test_wifi.c
+++ b/components/esp_wifi/test/test_wifi.c
@@ -1,8 +1,16 @@
 /*
- Tests for the Wi-Fi
-*/
+ * SPDX-FileCopyrightText: 2021 Espressif Systems (Shanghai) CO LTD
+ *
+ * SPDX-License-Identifier: Unlicense OR CC0-1.0
+ *
+ * This test code is in the Public Domain (or CC0 licensed, at your option.)
+ *
+ * Unless required by applicable law or agreed to in writing, this
+ * software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+ * CONDITIONS OF ANY KIND, either express or implied.
+ */
+
 #include "string.h"
-#include "esp_system.h"
 #include "unity.h"
 #include "esp_system.h"
 #include "esp_event.h"
@@ -11,6 +19,7 @@
 #include "esp_log.h"
 #include "nvs_flash.h"
 #include "test_utils.h"
+#include "memory_checks.h"
 #include "freertos/task.h"
 #include "freertos/event_groups.h"
 
@@ -190,7 +199,7 @@ static void start_wifi_as_softap(void)
     event_init();
 
     // can't deinit event loop, need to reset leak check
-    unity_reset_leak_checks();
+    test_utils_record_free_mem();
 
     if (wifi_events == NULL) {
         wifi_events = xEventGroupCreate();
@@ -212,7 +221,7 @@ static void start_wifi_as_sta(void)
     event_init();
 
     // can't deinit event loop, need to reset leak check
-    unity_reset_leak_checks();
+    test_utils_record_free_mem();
 
     if (wifi_events == NULL) {
         wifi_events = xEventGroupCreate();
diff --git a/components/mqtt/test/test_mqtt.c b/components/mqtt/test/test_mqtt.c
index ede3b31a79..e103b59d20 100644
--- a/components/mqtt/test/test_mqtt.c
+++ b/components/mqtt/test/test_mqtt.c
@@ -1,8 +1,21 @@
+/*
+ * SPDX-FileCopyrightText: 2021 Espressif Systems (Shanghai) CO LTD
+ *
+ * SPDX-License-Identifier: Unlicense OR CC0-1.0
+ *
+ * This test code is in the Public Domain (or CC0 licensed, at your option.)
+ *
+ * Unless required by applicable law or agreed to in writing, this
+ * software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+ * CONDITIONS OF ANY KIND, either express or implied.
+ */
+
 #include <sys/time.h>
 #include "freertos/FreeRTOS.h"
 #include "freertos/event_groups.h"
 #include "unity.h"
 #include "test_utils.h"
+#include "memory_checks.h"
 #include "mqtt_client.h"
 #include "nvs_flash.h"
 #include "esp_ota_ops.h"
@@ -17,7 +30,7 @@ static void test_leak_setup(const char * file, long line)
     gettimeofday(&te, NULL); // get current time
     esp_read_mac(mac, ESP_MAC_WIFI_STA);
     printf("%s:%ld: time=%ld.%lds, mac:" MACSTR "\n", file, line, te.tv_sec, te.tv_usec, MAC2STR(mac));
-    unity_reset_leak_checks();
+    test_utils_record_free_mem();
 }
 
 TEST_CASE("mqtt init with invalid url", "[mqtt][leaks=0]")
diff --git a/components/wpa_supplicant/test/test_offchannel.c b/components/wpa_supplicant/test/test_offchannel.c
index 72f2723bdc..c323beb6c1 100644
--- a/components/wpa_supplicant/test/test_offchannel.c
+++ b/components/wpa_supplicant/test/test_offchannel.c
@@ -1,16 +1,14 @@
-// Copyright 2020 Espressif Systems (Shanghai) PTE LTD
-//
-// 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.
+/*
+ * SPDX-FileCopyrightText: 2020-2021 Espressif Systems (Shanghai) CO LTD
+ *
+ * SPDX-License-Identifier: Unlicense OR CC0-1.0
+ *
+ * This test code is in the Public Domain (or CC0 licensed, at your option.)
+ *
+ * Unless required by applicable law or agreed to in writing, this
+ * software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+ * CONDITIONS OF ANY KIND, either express or implied.
+ */
 
 #include "string.h"
 #include "esp_system.h"
@@ -23,6 +21,7 @@
 #include "../esp_supplicant/src/esp_wifi_driver.h"
 #include "esp_log.h"
 #include "test_utils.h"
+#include "memory_checks.h"
 #include "freertos/event_groups.h"
 
 #define WIFI_START_EVENT        0x00000001
@@ -87,7 +86,7 @@ static void start_wifi_as_sta(void)
     event_init();
 
     // can't deinit event loop, need to reset leak check
-    unity_reset_leak_checks();
+    test_utils_record_free_mem();
 
     if (wifi_event == NULL) {
         wifi_event = xEventGroupCreate();
diff --git a/tools/ci/check_copyright_ignore.txt b/tools/ci/check_copyright_ignore.txt
index d5c845a7b4..273ea8adc3 100644
--- a/tools/ci/check_copyright_ignore.txt
+++ b/tools/ci/check_copyright_ignore.txt
@@ -823,7 +823,6 @@ components/esp_timer/test/test_esp_timer_light_sleep.c
 components/esp_timer/test/test_ets_timer.c
 components/esp_websocket_client/esp_websocket_client.c
 components/esp_websocket_client/include/esp_websocket_client.h
-components/esp_websocket_client/test/test_websocket_client.c
 components/esp_wifi/include/esp_coexist_adapter.h
 components/esp_wifi/include/esp_mesh.h
 components/esp_wifi/include/esp_mesh_internal.h
@@ -839,7 +838,6 @@ components/esp_wifi/src/lib_printf.c
 components/esp_wifi/src/mesh_event.c
 components/esp_wifi/src/smartconfig.c
 components/esp_wifi/src/smartconfig_ack.c
-components/esp_wifi/test/test_wifi.c
 components/esp_wifi/test/test_wifi_init.c
 components/espcoredump/corefile/__init__.py
 components/espcoredump/corefile/_parse_soc_header.py
@@ -1426,7 +1424,6 @@ components/mqtt/host_test/mocks/include/freertos/FreeRTOSConfig.h
 components/mqtt/host_test/mocks/include/freertos/portmacro.h
 components/mqtt/host_test/mocks/include/machine/endian.h
 components/mqtt/host_test/mocks/include/sys/queue.h
-components/mqtt/test/test_mqtt.c
 components/newlib/abort.c
 components/newlib/assert.c
 components/newlib/heap.c
@@ -2501,7 +2498,6 @@ components/wpa_supplicant/src/wps/wps_registrar.c
 components/wpa_supplicant/src/wps/wps_validate.c
 components/wpa_supplicant/test/test_crypto.c
 components/wpa_supplicant/test/test_dpp.c
-components/wpa_supplicant/test/test_offchannel.c
 components/wpa_supplicant/test/test_sae.c
 components/xtensa/eri.c
 components/xtensa/esp32/include/xtensa/config/core-isa.h
@@ -3318,13 +3314,10 @@ tools/test_mkdfu/test_mkdfu.py
 tools/test_mkuf2/test_mkuf2.py
 tools/unit-test-app/components/test_utils/ccomp_timer.c
 tools/unit-test-app/components/test_utils/include/ccomp_timer.h
-tools/unit-test-app/components/test_utils/include/test_utils.h
 tools/unit-test-app/components/test_utils/private_include/ccomp_timer_impl.h
 tools/unit-test-app/components/test_utils/test/ccomp_timer_test_api.c
 tools/unit-test-app/components/test_utils/test/ccomp_timer_test_data.c
 tools/unit-test-app/components/test_utils/test/ccomp_timer_test_inst.c
-tools/unit-test-app/components/test_utils/test_runner.c
-tools/unit-test-app/components/test_utils/test_utils.c
 tools/unit-test-app/idf_ext.py
 tools/unit-test-app/main/app_main.c
 tools/unit-test-app/tools/CreateSectionTable.py
diff --git a/tools/unit-test-app/components/test_utils/CMakeLists.txt b/tools/unit-test-app/components/test_utils/CMakeLists.txt
index b72d040dab..1fd457f492 100644
--- a/tools/unit-test-app/components/test_utils/CMakeLists.txt
+++ b/tools/unit-test-app/components/test_utils/CMakeLists.txt
@@ -1,4 +1,5 @@
 set(srcs "ccomp_timer.c"
+         "memory_checks.c"
          "test_runner.c"
          "test_utils.c")
 
diff --git a/tools/unit-test-app/components/test_utils/include/memory_checks.h b/tools/unit-test-app/components/test_utils/include/memory_checks.h
new file mode 100644
index 0000000000..636757e246
--- /dev/null
+++ b/tools/unit-test-app/components/test_utils/include/memory_checks.h
@@ -0,0 +1,91 @@
+/*
+ * SPDX-FileCopyrightText: 2015-2021 Espressif Systems (Shanghai) CO LTD
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <stddef.h>
+#include "esp_err.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * @brief Leak for components
+ */
+typedef enum {
+    ESP_COMP_LEAK_GENERAL = 0,  /**< Leak by default */
+    ESP_COMP_LEAK_LWIP,         /**< Leak for LWIP */
+    ESP_COMP_LEAK_NVS,          /**< Leak for NVS */
+    ESP_COMP_LEAK_ALL,          /**< Use for getting the summary leak level */
+} esp_comp_leak_t;
+
+/**
+ * @brief Type of a leak threshold
+ */
+typedef enum {
+    ESP_LEAK_TYPE_WARNING = 0,  /**< Warning level of leak */
+    ESP_LEAK_TYPE_CRITICAL,     /**< Critical level of leak */
+    ESP_LEAK_TYPE_MAX,          /**< Max number of leak levels for all components/levels */
+} esp_type_leak_t;
+
+/**
+ * @brief Adjust the memory leak thresholds for unit tests.
+ *
+ * Usually, unit tests will check if memory is leaked. Some functionality used by unit tests may unavoidably
+ * leak memory. This is why there is a default threshold for memory leaks (currently 1024 bytes).
+ * Within this range, the number of bytes leaked will be visually reported on the terminal, but no test failure will
+ * be triggered. Any memory leak above the default threshold will trigger a unit test failure.
+ * This function allows to adjust that memory leak threshold.
+ *
+ * @param leak_level Maximum allowed memory leak which will not trigger a unit test failure.
+ * @param type_of_leak There are two types of leak thresholds: critical and warning. Only the
+ *                      critical threshold will trigger a unit test failure if exceeded.
+ * @param component Thresholds can be set in general or for specific components. Note that this argument
+ *                      is not checked.
+ *
+ * @return ESP_OK on success, ESP_INVALID_ARG if type_of_leak is invalid. \c component is unchecked.
+ */
+esp_err_t test_utils_set_leak_level(size_t leak_level, esp_type_leak_t type_of_leak, esp_comp_leak_t component);
+
+/**
+ * @brief Return the memory leak thresholds for unit tests for a leak type and component.
+ *
+ * For more information, see \c test_utils_set_leak_level above.
+ *
+ * @param type_of_leak Warning or Critical
+ * @param component The component for which to return the leak threshold.
+ */
+size_t test_utils_get_leak_level(esp_type_leak_t type_of_leak, esp_comp_leak_t component);
+
+/**
+ * @brief Start/Restart memory leak checking.
+ *
+ * Records the current free memory values at time of calling. After the test case, it may be checked with
+ * \c test_utils_finish_and_evaluate_leaks.
+ *
+ * If this function is called repeatedly, only the free memory values at the last time of calling will prevail
+ * as reference.
+ */
+void test_utils_record_free_mem(void);
+
+/**
+ * @brief Evaluate memory leak checking according to the provided thresholds.
+ *
+ * If the current memory leak level (counted from the last time calling \c test_utils_record_free_mem() ) exceeds
+ * \c critical_threshold, a unit test failure will be triggered. If it exceeds only the warning,
+ * a warning message will be issued.
+ */
+void test_utils_finish_and_evaluate_leaks(size_t warn_threshold, size_t critical_threshold);
+
+/**
+ * @brief Helper function to setup and initialize heap tracing.
+ */
+void setup_heap_record(void);
+
+#ifdef __cplusplus
+}
+#endif
diff --git a/tools/unit-test-app/components/test_utils/include/test_utils.h b/tools/unit-test-app/components/test_utils/include/test_utils.h
index c9a5122a81..a739eb1fd8 100644
--- a/tools/unit-test-app/components/test_utils/include/test_utils.h
+++ b/tools/unit-test-app/components/test_utils/include/test_utils.h
@@ -1,16 +1,9 @@
-// Copyright 2015-2018 Espressif Systems (Shanghai) PTE LTD
-//
-// 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.
+/*
+ * SPDX-FileCopyrightText: 2015-2021 Espressif Systems (Shanghai) CO LTD
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
 #pragma once
 
 // Utilities for esp-idf unit tests
@@ -118,18 +111,6 @@ uint64_t ref_clock_get(void);
  */
 void test_main(void);
 
-/**
- * @brief Reset automatic leak checking which happens in unit tests.
- *
- * Updates recorded "before" free memory values to the free memory values
- * at time of calling. Resets leak checker if tracing is enabled in
- * config.
- *
- * This can be called if a test case does something which allocates
- * memory on first use, for example.
- *
- * @note Use with care as this can mask real memory leak problems.
- */
 void unity_reset_leak_checks(void);
 
 
@@ -232,47 +213,6 @@ static inline void unity_send_signal(const char* signal_name)
  */
 bool unity_util_convert_mac_from_string(const char* mac_str, uint8_t *mac_addr);
 
-/**
- * @brief Leak for components
- */
-typedef enum {
-    COMP_LEAK_GENERAL = 0,  /**< Leak by default */
-    COMP_LEAK_LWIP,         /**< Leak for LWIP */
-    COMP_LEAK_NVS,          /**< Leak for NVS */
-    COMP_LEAK_ALL,          /**< Use for getting the summary leak level */
-} esp_comp_leak_t;
-
-/**
- * @brief Type of leak
- */
-typedef enum {
-    TYPE_LEAK_WARNING = 0,  /**< Warning level of leak */
-    TYPE_LEAK_CRITICAL,     /**< Critical level of leak */
-    TYPE_LEAK_MAX,          /**< Max number of leak levels */
-} esp_type_leak_t;
-
-/**
- * @brief Set a leak level for the required type and component.
- *
- * @param[in] leak_level Level of leak
- * @param[in] type       Type of leak
- * @param[in] component  Name of component
- *
- * return ESP_OK: Successful.
- *        ESP_ERR_INVALID_ARG: Invalid argument.
- */
-esp_err_t test_utils_set_leak_level(size_t leak_level, esp_type_leak_t type, esp_comp_leak_t component);
-
-/**
- * @brief Get a leak level for the required type and component.
- *
- * @param[in] type       Type of leak.
- * @param[in] component  Name of component. If COMP_LEAK_ALL, then the level will be summarized for all components.
- * return Leak level
- */
-size_t test_utils_get_leak_level(esp_type_leak_t type, esp_comp_leak_t component);
-
-
 typedef struct test_utils_exhaust_memory_record_s *test_utils_exhaust_memory_rec;
 
 /**
diff --git a/tools/unit-test-app/components/test_utils/memory_checks.c b/tools/unit-test-app/components/test_utils/memory_checks.c
new file mode 100644
index 0000000000..f126ea9c12
--- /dev/null
+++ b/tools/unit-test-app/components/test_utils/memory_checks.c
@@ -0,0 +1,97 @@
+/*
+ * SPDX-FileCopyrightText: 2015-2021 Espressif Systems (Shanghai) CO LTD
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <stdio.h>
+#include "esp_heap_caps.h"
+#include "unity.h"
+#include "memory_checks.h"
+
+static size_t before_free_8bit;
+static size_t before_free_32bit;
+
+static size_t test_unity_leak_level[ESP_LEAK_TYPE_MAX][ESP_COMP_LEAK_ALL] = { 0 };
+
+esp_err_t test_utils_set_leak_level(size_t leak_level, esp_type_leak_t type_of_leak, esp_comp_leak_t component)
+{
+    if (type_of_leak >= ESP_LEAK_TYPE_MAX || component >= ESP_COMP_LEAK_ALL) {
+        return ESP_ERR_INVALID_ARG;
+    }
+    test_unity_leak_level[type_of_leak][component] = leak_level;
+    return ESP_OK;
+}
+
+size_t test_utils_get_leak_level(esp_type_leak_t type_of_leak, esp_comp_leak_t component)
+{
+    size_t leak_level = 0;
+    if (type_of_leak >= ESP_LEAK_TYPE_MAX || component > ESP_COMP_LEAK_ALL) {
+        leak_level = 0;
+    } else {
+        if (component == ESP_COMP_LEAK_ALL) {
+            for (int comp = 0; comp < ESP_COMP_LEAK_ALL; ++comp) {
+                leak_level += test_unity_leak_level[type_of_leak][comp];
+            }
+        } else {
+            leak_level = test_unity_leak_level[type_of_leak][component];
+        }
+    }
+    return leak_level;
+}
+
+void test_utils_record_free_mem(void)
+{
+    before_free_8bit = heap_caps_get_free_size(MALLOC_CAP_8BIT);
+    before_free_32bit = heap_caps_get_free_size(MALLOC_CAP_32BIT);
+}
+
+void setup_heap_record(void)
+{
+#ifdef CONFIG_HEAP_TRACING
+    const size_t num_heap_records = 80;
+    static heap_trace_record_t *record_buffer;
+    if (!record_buffer) {
+        record_buffer = malloc(sizeof(heap_trace_record_t) * num_heap_records);
+        assert(record_buffer);
+        heap_trace_init_standalone(record_buffer, num_heap_records);
+    }
+#endif
+}
+
+static void check_leak(size_t before_free,
+        size_t after_free,
+        const char *type,
+        size_t warn_threshold,
+        size_t critical_threshold)
+{
+    int free_delta = (int)after_free - (int)before_free;
+    printf("MALLOC_CAP_%s usage: Free memory delta: %d Leak threshold: -%u \n",
+           type,
+           free_delta,
+           critical_threshold);
+
+    if (free_delta > 0) {
+        return; // free memory went up somehow
+    }
+
+    size_t leaked = (size_t)(free_delta * -1);
+    if (leaked <= warn_threshold) {
+        return;
+    }
+
+    printf("MALLOC_CAP_%s %s leak: Before %u bytes free, After %u bytes free (delta %u)\n",
+           type,
+           leaked <= critical_threshold ? "potential" : "critical",
+           before_free, after_free, leaked);
+    fflush(stdout);
+    TEST_ASSERT_MESSAGE(leaked <= critical_threshold, "The test leaked too much memory");
+}
+
+void test_utils_finish_and_evaluate_leaks(size_t warn_threshold, size_t critical_threshold)
+{
+    size_t after_free_8bit = heap_caps_get_free_size(MALLOC_CAP_8BIT);
+    size_t after_free_32bit = heap_caps_get_free_size(MALLOC_CAP_32BIT);
+    check_leak(before_free_8bit, after_free_8bit, "8BIT", warn_threshold, critical_threshold);
+    check_leak(before_free_32bit, after_free_32bit, "32BIT", warn_threshold, critical_threshold);
+}
diff --git a/tools/unit-test-app/components/test_utils/test_runner.c b/tools/unit-test-app/components/test_utils/test_runner.c
index 3584a9068c..f9e1422361 100644
--- a/tools/unit-test-app/components/test_utils/test_runner.c
+++ b/tools/unit-test-app/components/test_utils/test_runner.c
@@ -1,16 +1,8 @@
-// Copyright 2016-2018 Espressif Systems (Shanghai) PTE LTD
-//
-// 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.
+/*
+ * SPDX-FileCopyrightText: 2015-2021 Espressif Systems (Shanghai) CO LTD
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
 
 #include <stdio.h>
 #include "string.h"
@@ -21,17 +13,12 @@
 #include "unity_test_runner.h"
 #include "test_utils.h"
 #include "esp_newlib.h"
+#include "memory_checks.h"
 
 #ifdef CONFIG_HEAP_TRACING
 #include "esp_heap_trace.h"
 #endif
 
-static size_t before_free_8bit;
-static size_t before_free_32bit;
-
-static size_t warn_leak_threshold;
-static size_t critical_leak_threshold;
-
 static void unity_task(void *pvParameters)
 {
     vTaskDelay(2); /* Delay a bit to let the main task be deleted */
@@ -46,28 +33,12 @@ void test_main(void)
                             UNITY_FREERTOS_PRIORITY, NULL, UNITY_FREERTOS_CPU);
 }
 
-void unity_reset_leak_checks(void)
-{
-    before_free_8bit = heap_caps_get_free_size(MALLOC_CAP_8BIT);
-    before_free_32bit = heap_caps_get_free_size(MALLOC_CAP_32BIT);
-
-#ifdef CONFIG_HEAP_TRACING
-    heap_trace_start(HEAP_TRACE_LEAKS);
-#endif
-}
-
 /* setUp runs before every test */
 void setUp(void)
 {
 // If heap tracing is enabled in kconfig, leak trace the test
 #ifdef CONFIG_HEAP_TRACING
-    const size_t num_heap_records = 80;
-    static heap_trace_record_t *record_buffer;
-    if (!record_buffer) {
-        record_buffer = malloc(sizeof(heap_trace_record_t) * num_heap_records);
-        assert(record_buffer);
-        heap_trace_init_standalone(record_buffer, num_heap_records);
-    }
+    setup_heap_record();
 #endif
 
     printf("%s", ""); /* sneakily lazy-allocate the reent structure for this test task */
@@ -81,55 +52,42 @@ void setUp(void)
     get_test_data_partition();  /* allocate persistent partition table structures */
 #endif // CONFIG_APP_BUILD_USE_FLASH_SECTIONS
 
-    unity_reset_leak_checks();
-    test_utils_set_leak_level(CONFIG_UNITY_CRITICAL_LEAK_LEVEL_GENERAL, TYPE_LEAK_CRITICAL, COMP_LEAK_GENERAL);
-    test_utils_set_leak_level(CONFIG_UNITY_WARN_LEAK_LEVEL_GENERAL, TYPE_LEAK_WARNING, COMP_LEAK_GENERAL);
+#ifdef CONFIG_HEAP_TRACING
+    heap_trace_start(HEAP_TRACE_LEAKS);
+#endif
+    test_utils_record_free_mem();
+    test_utils_set_leak_level(CONFIG_UNITY_CRITICAL_LEAK_LEVEL_GENERAL, ESP_LEAK_TYPE_CRITICAL, ESP_COMP_LEAK_GENERAL);
+    test_utils_set_leak_level(CONFIG_UNITY_WARN_LEAK_LEVEL_GENERAL, ESP_LEAK_TYPE_WARNING, ESP_COMP_LEAK_GENERAL);
 }
 
-static void check_leak(size_t before_free, size_t after_free, const char *type)
+typedef enum {
+    NO_LEAK_CHECK,
+    DEFAULT_LEAK_CHECK,
+    SPECIAL_LEAK_CHECK
+} leak_check_type_t;
+
+/**
+ * It is possible to specify the maximum allowed memory leak level directly in the test case
+ * or disable leak checking for a test case.
+ * This function checks if this is the case and return the appropriate return value.
+ * If a custom leak level has been specified, that custom threshold is written to the value pointed by threshold.
+ */
+static leak_check_type_t leak_check_required(size_t *threshold)
 {
-    int free_delta = (int)after_free - (int)before_free;
-    printf("MALLOC_CAP_%s usage: Free memory delta: %d Leak threshold: -%u \n",
-           type,
-           free_delta,
-           critical_leak_threshold);
-
-    if (free_delta > 0) {
-        return; // free memory went up somehow
-    }
-
-    size_t leaked = (size_t)(free_delta * -1);
-    if (leaked <= warn_leak_threshold) {
-        return;
-    }
-
-    printf("MALLOC_CAP_%s %s leak: Before %u bytes free, After %u bytes free (delta %u)\n",
-           type,
-           leaked <= critical_leak_threshold ? "potential" : "critical",
-           before_free, after_free, leaked);
-    fflush(stdout);
-    TEST_ASSERT_MESSAGE(leaked <= critical_leak_threshold, "The test leaked too much memory");
-}
-
-static bool leak_check_required(void)
-{
-    warn_leak_threshold = test_utils_get_leak_level(TYPE_LEAK_WARNING, COMP_LEAK_ALL);
-    critical_leak_threshold = test_utils_get_leak_level(TYPE_LEAK_CRITICAL, COMP_LEAK_ALL);
     if (Unity.CurrentDetail1 != NULL) {
         const char *leaks = "[leaks";
         const int len_leaks = strlen(leaks);
         const char *sub_leaks = strstr(Unity.CurrentDetail1, leaks);
         if (sub_leaks != NULL) {
             if (sub_leaks[len_leaks] == ']') {
-                return false;
+                return NO_LEAK_CHECK;
             } else if (sub_leaks[len_leaks] == '=') {
-                critical_leak_threshold = strtol(&sub_leaks[len_leaks + 1], NULL, 10);
-                warn_leak_threshold = critical_leak_threshold;
-                return true;
+                *threshold = strtol(&sub_leaks[len_leaks + 1], NULL, 10);
+                return SPECIAL_LEAK_CHECK;
             }
         }
     }
-    return true;
+    return DEFAULT_LEAK_CHECK;
 }
 
 /* tearDown runs after every test */
@@ -154,11 +112,23 @@ void tearDown(void)
     heap_trace_dump();
 #endif
 
-    if (leak_check_required()) {
-          size_t after_free_8bit = heap_caps_get_free_size(MALLOC_CAP_8BIT);
-          size_t after_free_32bit = heap_caps_get_free_size(MALLOC_CAP_32BIT);
-          check_leak(before_free_8bit, after_free_8bit, "8BIT");
-          check_leak(before_free_32bit, after_free_32bit, "32BIT");
+    size_t leak_threshold_critical = 0;
+    size_t leak_threshold_warning = 0;
+    leak_check_type_t check_type = leak_check_required(&leak_threshold_critical);
+
+    // In the "special case", only one level can be passed directly from the test case.
+    // Hence, we set both warning and critical leak levels to that same value here
+    leak_threshold_warning = leak_threshold_critical;
+
+    if (check_type == NO_LEAK_CHECK) {
+        // do not check
+    } else if (check_type == SPECIAL_LEAK_CHECK) {
+        test_utils_finish_and_evaluate_leaks(leak_threshold_warning, leak_threshold_critical);
+    } else if (check_type == DEFAULT_LEAK_CHECK) {
+        test_utils_finish_and_evaluate_leaks(test_utils_get_leak_level(ESP_LEAK_TYPE_WARNING, ESP_COMP_LEAK_ALL),
+                test_utils_get_leak_level(ESP_LEAK_TYPE_CRITICAL, ESP_COMP_LEAK_ALL));
+    } else {
+        assert(false); // coding error
     }
 
     Unity.TestFile = real_testfile; // go back to the real filename
diff --git a/tools/unit-test-app/components/test_utils/test_utils.c b/tools/unit-test-app/components/test_utils/test_utils.c
index 100bdf9533..1d4b6cf315 100644
--- a/tools/unit-test-app/components/test_utils/test_utils.c
+++ b/tools/unit-test-app/components/test_utils/test_utils.c
@@ -1,16 +1,8 @@
-// Copyright 2015-2016 Espressif Systems (Shanghai) PTE LTD
-//
-// 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.
+/*
+ * SPDX-FileCopyrightText: 2015-2021 Espressif Systems (Shanghai) CO LTD
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
 
 #include <string.h>
 #include "unity.h"
@@ -20,6 +12,7 @@
 #include "esp_netif.h"
 #include "lwip/sockets.h"
 #include "sdkconfig.h"
+#include "memory_checks.h"
 #if !CONFIG_FREERTOS_UNICORE
 #include "esp_ipc.h"
 #include "esp_freertos_hooks.h"
@@ -57,9 +50,9 @@ void test_case_uses_tcpip(void)
     printf("Note: esp_netif_init() has been called. Until next reset, TCP/IP task will periodicially allocate memory and consume CPU time.\n");
 
     // Reset the leak checker as LWIP allocates a lot of memory on first run
-    unity_reset_leak_checks();
-    test_utils_set_leak_level(0, TYPE_LEAK_CRITICAL, COMP_LEAK_GENERAL);
-    test_utils_set_leak_level(CONFIG_UNITY_CRITICAL_LEAK_LEVEL_LWIP, TYPE_LEAK_CRITICAL, COMP_LEAK_LWIP);
+    test_utils_record_free_mem();
+    test_utils_set_leak_level(0, ESP_LEAK_TYPE_CRITICAL, ESP_COMP_LEAK_GENERAL);
+    test_utils_set_leak_level(CONFIG_UNITY_CRITICAL_LEAK_LEVEL_LWIP, ESP_LEAK_TYPE_CRITICAL, ESP_COMP_LEAK_LWIP);
 }
 
 // wait user to send "Enter" key or input parameter
@@ -119,34 +112,6 @@ bool unity_util_convert_mac_from_string(const char* mac_str, uint8_t *mac_addr)
     return true;
 }
 
-static size_t test_unity_leak_level[TYPE_LEAK_MAX][COMP_LEAK_ALL] = { 0 };
-
-esp_err_t test_utils_set_leak_level(size_t leak_level, esp_type_leak_t type_of_leak, esp_comp_leak_t component)
-{
-    if (type_of_leak >= TYPE_LEAK_MAX || component >= COMP_LEAK_ALL) {
-        return ESP_ERR_INVALID_ARG;
-    }
-    test_unity_leak_level[type_of_leak][component] = leak_level;
-    return ESP_OK;
-}
-
-size_t test_utils_get_leak_level(esp_type_leak_t type_of_leak, esp_comp_leak_t component)
-{
-    size_t leak_level = 0;
-    if (type_of_leak >= TYPE_LEAK_MAX || component > COMP_LEAK_ALL) {
-        leak_level = 0;
-    } else {
-        if (component == COMP_LEAK_ALL) {
-            for (int comp = 0; comp < COMP_LEAK_ALL; ++comp) {
-                leak_level += test_unity_leak_level[type_of_leak][comp];
-            }
-        } else {
-            leak_level = test_unity_leak_level[type_of_leak][component];
-        }
-    }
-    return leak_level;
-}
-
 #define EXHAUST_MEMORY_ENTRIES 100
 
 struct test_utils_exhaust_memory_record_s {