diff --git a/include/test/optee.h b/include/test/optee.h
index f4255b3..0a548a5 100644
--- a/include/test/optee.h
+++ b/include/test/optee.h
@@ -8,7 +8,9 @@
 
 #include <test/test.h>
 
-/* Declare a new environment test */
+/* Declare a new optee test */
 #define OPTEE_TEST(_name, _flags)	UNIT_TEST(_name, _flags, optee)
+#define OPTEE_TEST_INIT(_name, _flags)	UNIT_TEST_INIT(_name, _flags, optee)
+#define OPTEE_TEST_UNINIT(_name, _flags) UNIT_TEST_UNINIT(_name, _flags, optee)
 
 #endif /* __TEST_OPTEE_H__ */
diff --git a/test/optee/cmd_ut_optee.c b/test/optee/cmd_ut_optee.c
index fc66747..792de30 100644
--- a/test/optee/cmd_ut_optee.c
+++ b/test/optee/cmd_ut_optee.c
@@ -26,6 +26,41 @@
 static void *fdt;
 static bool expect_success;
 
+static int optee_test_init(struct unit_test_state *uts)
+{
+	void *fdt_optee = &__dtb_test_optee_optee_begin;
+	void *fdt_no_optee = &__dtb_test_optee_no_optee_begin;
+	void *fdt_base = &__dtb_test_optee_base_begin;
+	int ret = -ENOMEM;
+
+	ut_assertok(fdt_check_header(fdt_base));
+	ut_assertok(fdt_check_header(fdt_optee));
+	ut_assertok(fdt_check_header(fdt_no_optee));
+
+	fdt = malloc(FDT_COPY_SIZE);
+	if (!fdt)
+		return ret;
+
+	/*
+	 * Resize the FDT to 4k so that we have room to operate on
+	 *
+	 * (and relocate it since the memory might be mapped
+	 * read-only)
+	 */
+	ut_assertok(fdt_open_into(fdt_base, fdt, FDT_COPY_SIZE));
+
+	return 0;
+}
+OPTEE_TEST_INIT(optee_test_init, 0);
+
+static int optee_test_uninit(struct unit_test_state *uts)
+{
+	free(fdt);
+
+	return 0;
+}
+OPTEE_TEST_UNINIT(optee_test_uninit, 0);
+
 static int optee_fdt_firmware(struct unit_test_state *uts)
 {
 	const void *prop;
@@ -101,26 +136,6 @@
 	void *fdt_base = &__dtb_test_optee_base_begin;
 	int ret = -ENOMEM;
 
-	uts = calloc(1, sizeof(*uts));
-	if (!uts)
-		return -ENOMEM;
-
-	ut_assertok(fdt_check_header(fdt_base));
-	ut_assertok(fdt_check_header(fdt_optee));
-	ut_assertok(fdt_check_header(fdt_no_optee));
-
-	fdt = malloc(FDT_COPY_SIZE);
-	if (!fdt)
-		return ret;
-
-	/*
-	 * Resize the FDT to 4k so that we have room to operate on
-	 *
-	 * (and relocate it since the memory might be mapped
-	 * read-only)
-	 */
-	ut_assertok(fdt_open_into(fdt_base, fdt, FDT_COPY_SIZE));
-
 	/*
 	 * (1) Try to copy optee nodes from empty dt.
 	 * This should still run successfully.