diff --git a/persistence-modules/spring-boot-persistence-2/src/main/java/com/baeldung/db/indexing/Image.java b/persistence-modules/spring-boot-persistence-2/src/main/java/com/baeldung/db/indexing/Image.java
index 5348788c12..e3fcf53f81 100644
--- a/persistence-modules/spring-boot-persistence-2/src/main/java/com/baeldung/db/indexing/Image.java
+++ b/persistence-modules/spring-boot-persistence-2/src/main/java/com/baeldung/db/indexing/Image.java
@@ -22,6 +22,10 @@ class Image {
     public Image() {
     }
 
+    public Image(Long id) {
+        this.id = id;
+    }
+
     public Image(String name, String location) {
         this.name = name;
         this.location = location;
diff --git a/persistence-modules/spring-boot-persistence-2/src/test/java/com/baeldung/db/indexing/FileSystemImageIntegrationTest.java b/persistence-modules/spring-boot-persistence-2/src/test/java/com/baeldung/db/indexing/FileSystemImageIntegrationTest.java
index 3082f16a78..83f5bae095 100644
--- a/persistence-modules/spring-boot-persistence-2/src/test/java/com/baeldung/db/indexing/FileSystemImageIntegrationTest.java
+++ b/persistence-modules/spring-boot-persistence-2/src/test/java/com/baeldung/db/indexing/FileSystemImageIntegrationTest.java
@@ -35,8 +35,12 @@ class FileSystemImageIntegrationTest {
         ClassLoader classLoader = ClassLoader.getSystemClassLoader();
         InputStream image = classLoader.getResourceAsStream("baeldung.jpeg");
 
+        MockMultipartFile jpegImage = new MockMultipartFile("image", "baeldung", MediaType.TEXT_PLAIN_VALUE, image);
         MockMultipartHttpServletRequestBuilder multipartRequest = MockMvcRequestBuilders.multipart("/file-system/image")
-            .file(new MockMultipartFile("image", "baeldung", MediaType.TEXT_PLAIN_VALUE, image));
+            .file(jpegImage);
+
+        given(fileLocationService.save(jpegImage.getBytes(), "baeldung"))
+            .willReturn(1L);
 
         MvcResult result = mockMvc.perform(multipartRequest)
             .andExpect(status().isOk())
diff --git a/persistence-modules/spring-boot-persistence-2/src/test/java/com/baeldung/db/indexing/ImageIntegrationTest.java b/persistence-modules/spring-boot-persistence-2/src/test/java/com/baeldung/db/indexing/ImageIntegrationTest.java
index f7d4ecf129..e38e0a21a9 100644
--- a/persistence-modules/spring-boot-persistence-2/src/test/java/com/baeldung/db/indexing/ImageIntegrationTest.java
+++ b/persistence-modules/spring-boot-persistence-2/src/test/java/com/baeldung/db/indexing/ImageIntegrationTest.java
@@ -1,6 +1,7 @@
 package com.baeldung.db.indexing;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.BDDMockito.given;
 import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status;
 
@@ -34,13 +35,10 @@ class ImageIntegrationTest {
 
     @Test
     void givenBaeldungJpegImage_whenUploadIt_thenReturnItsId() throws Exception {
-        ClassLoader classLoader = ClassLoader.getSystemClassLoader();
-        InputStream image = classLoader.getResourceAsStream("baeldung.jpeg");
+        given(imageRepository.save(any()))
+            .willReturn(new Image(1L));
 
-        MockMultipartHttpServletRequestBuilder multipartRequest = MockMvcRequestBuilders.multipart("/image")
-            .file(new MockMultipartFile("image", "baeldung", MediaType.TEXT_PLAIN_VALUE, image));
-
-        MvcResult result = mockMvc.perform(multipartRequest)
+        MvcResult result = mockMvc.perform(createUploadImageRequest())
             .andExpect(status().isOk())
             .andReturn();
 
@@ -60,6 +58,14 @@ class ImageIntegrationTest {
             .andExpect(status().isOk());
     }
 
+    private MockMultipartHttpServletRequestBuilder createUploadImageRequest() throws IOException {
+        ClassLoader classLoader = ClassLoader.getSystemClassLoader();
+        InputStream image = classLoader.getResourceAsStream("baeldung.jpeg");
+
+        return MockMvcRequestBuilders.multipart("/image")
+            .file(new MockMultipartFile("multipartImage", "baeldung", MediaType.TEXT_PLAIN_VALUE, image));
+    }
+
     private Image baeldungImage() throws IOException {
         ClassLoader classLoader = ClassLoader.getSystemClassLoader();