diff --git a/source/core/flash.c b/source/core/flash.c
index 2d37908..d1ab66b 100644
--- a/source/core/flash.c
+++ b/source/core/flash.c
@@ -125,6 +125,11 @@ size_t flash_write(struct flash_dev *dev, uint32_t addr, const void *data,
 #define flash_write __wrap_flash_write
 #endif
 
+#ifdef flash_write0
+#undef flash_write0
+#define flash_write0 __real_flash_write0
+#endif
+
 size_t flash_write0(struct flash_dev *dev, uint32_t addr, size_t len)
 {
 	uint8_t data[32];
@@ -146,6 +151,16 @@ size_t flash_write0(struct flash_dev *dev, uint32_t addr, size_t len)
 	return ret;
 }
 
+#ifdef flash_write0
+#undef flash_write0
+#define flash_write0 __wrap_flash_write0
+#endif
+
+#ifdef flash_copy
+#undef flash_copy
+#define flash_copy __real_flash_copy
+#endif
+
 size_t flash_copy(struct flash_dev *dev, uint32_t dst, uint32_t src, size_t len)
 {
 	if (!dev)
@@ -154,6 +169,11 @@ size_t flash_copy(struct flash_dev *dev, uint32_t dst, uint32_t src, size_t len)
 	return dev->ops->copy(dev, dst, src, len);
 }
 
+#ifdef flash_copy
+#undef flash_copy
+#define flash_copy __wrap_flash_copy
+#endif
+
 #ifdef flash_is_erased
 #undef flash_is_erased
 #define flash_is_erased __real_flash_is_erased
diff --git a/source/tests/Makefile b/source/tests/Makefile
index e9e7064..faa972e 100644
--- a/source/tests/Makefile
+++ b/source/tests/Makefile
@@ -21,6 +21,8 @@ test-obj-y += source/tests/ftl/ftl_is_mapped.o
 
 TEST_CFLAGS += -Dflash_read=__wrap_flash_read
 TEST_CFLAGS += -Dflash_write=__wrap_flash_write
+TEST_CFLAGS += -Dflash_write0=__wrap_flash_write0
+TEST_CFLAGS += -Dflash_copy=__wrap_flash_copy
 TEST_CFLAGS += -Dflash_is_erased=__wrap_flash_is_erased
 
 TEST_CFLAGS += -Dis_group_erased=__wrap_is_group_erased
diff --git a/source/tests/flash/mock.c b/source/tests/flash/mock.c
index 4aca2e9..5974f25 100644
--- a/source/tests/flash/mock.c
+++ b/source/tests/flash/mock.c
@@ -57,6 +57,25 @@ size_t __wrap_flash_write(struct flash_dev *dev, uint32_t addr,
 	return len;
 }
 
+size_t __wrap_flash_write0(struct flash_dev *dev, uint32_t addr, size_t len)
+{
+	check_expected(addr);
+	check_expected(len);
+
+	return mock_type(int);
+}
+
+size_t __wrap_flash_copy(struct flash_dev *dev, uint32_t dst, uint32_t src, size_t len)
+{
+	(void)dev;
+
+	check_expected(dst);
+	check_expected(src);
+	check_expected(len);
+
+	return mock_type(int);
+}
+
 int __wrap_flash_is_erased(struct flash_dev *dev, uint32_t addr, size_t len)
 {
 	(void)dev;