diff --git a/core-java-modules/core-java-io-2/src/main/java/com/baeldung/listfiles/ListFiles.java b/core-java-modules/core-java-io-2/src/main/java/com/baeldung/listfiles/ListFiles.java
index c108d9ed96..ccc4f8808e 100644
--- a/core-java-modules/core-java-io-2/src/main/java/com/baeldung/listfiles/ListFiles.java
+++ b/core-java-modules/core-java-io-2/src/main/java/com/baeldung/listfiles/ListFiles.java
@@ -2,12 +2,7 @@ package com.baeldung.listfiles;
 
 import java.io.File;
 import java.io.IOException;
-import java.nio.file.DirectoryStream;
-import java.nio.file.FileVisitResult;
-import java.nio.file.Files;
-import java.nio.file.Path;
-import java.nio.file.Paths;
-import java.nio.file.SimpleFileVisitor;
+import java.nio.file.*;
 import java.nio.file.attribute.BasicFileAttributes;
 import java.util.HashSet;
 import java.util.Set;
@@ -15,32 +10,31 @@ import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 public class ListFiles {
-    public static final int DEPTH = 1;
 
     public Set<String> listFilesUsingJavaIO(String dir) {
         return Stream.of(new File(dir).listFiles())
-            .filter(file -> !file.isDirectory())
-            .map(File::getName)
-            .collect(Collectors.toSet());
+                .filter(file -> !file.isDirectory())
+                .map(File::getName)
+                .collect(Collectors.toSet());
     }
 
     public Set<String> listFilesUsingFilesList(String dir) throws IOException {
         try (Stream<Path> stream = Files.list(Paths.get(dir))) {
             return stream
-                .filter(file -> !Files.isDirectory(file))
-                .map(Path::getFileName)
-                .map(Path::toString)
-                .collect(Collectors.toSet());
+                    .filter(file -> !Files.isDirectory(file))
+                    .map(Path::getFileName)
+                    .map(Path::toString)
+                    .collect(Collectors.toSet());
         }
     }
 
     public Set<String> listFilesUsingFileWalk(String dir, int depth) throws IOException {
         try (Stream<Path> stream = Files.walk(Paths.get(dir), depth)) {
             return stream
-                .filter(file -> !Files.isDirectory(file))
-                .map(Path::getFileName)
-                .map(Path::toString)
-                .collect(Collectors.toSet());
+                    .filter(file -> !Files.isDirectory(file))
+                    .map(Path::getFileName)
+                    .map(Path::toString)
+                    .collect(Collectors.toSet());
         }
     }
 
@@ -48,10 +42,9 @@ public class ListFiles {
         Set<String> fileList = new HashSet<>();
         Files.walkFileTree(Paths.get(dir), new SimpleFileVisitor<Path>() {
             @Override
-            public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
+            public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) {
                 if (!Files.isDirectory(file)) {
-                    fileList.add(file.getFileName()
-                        .toString());
+                    fileList.add(file.getFileName().toString());
                 }
                 return FileVisitResult.CONTINUE;
             }
@@ -60,16 +53,14 @@ public class ListFiles {
     }
 
     public Set<String> listFilesUsingDirectoryStream(String dir) throws IOException {
-        Set<String> fileList = new HashSet<>();
+        Set<String> fileSet = new HashSet<>();
         try (DirectoryStream<Path> stream = Files.newDirectoryStream(Paths.get(dir))) {
             for (Path path : stream) {
                 if (!Files.isDirectory(path)) {
-                    fileList.add(path.getFileName()
-                        .toString());
+                    fileSet.add(path.getFileName().toString());
                 }
             }
         }
-        return fileList;
+        return fileSet;
     }
-
 }
diff --git a/core-java-modules/core-java-io-2/src/test/java/com/baeldung/listfiles/ListFilesUnitTest.java b/core-java-modules/core-java-io-2/src/test/java/com/baeldung/listfiles/ListFilesUnitTest.java
index 5b4efc9e58..bdf96f3df8 100644
--- a/core-java-modules/core-java-io-2/src/test/java/com/baeldung/listfiles/ListFilesUnitTest.java
+++ b/core-java-modules/core-java-io-2/src/test/java/com/baeldung/listfiles/ListFilesUnitTest.java
@@ -1,21 +1,21 @@
 package com.baeldung.listfiles;
 
-import static org.junit.Assert.assertEquals;
-
-import java.io.IOException;
-import java.util.HashSet;
-import java.util.Set;
-
 import org.junit.Test;
 
-import com.baeldung.listfiles.ListFiles;
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
+
+import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
+import static org.junit.Assert.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 public class ListFilesUnitTest {
 
-    private ListFiles listFiles = new ListFiles();
-    private String DIRECTORY = "src/test/resources/listFilesUnitTestFolder";
-    private static final int DEPTH = 1;
-    private Set<String> EXPECTED_FILE_LIST = new HashSet<String>() {
+    private static final String VALID_DIRECTORY = "src/test/resources/listFilesUnitTestFolder";
+    private static final String INVALID_DIRECTORY = "src/test/resources/thisDirectoryDoesNotExist";
+    private static final Set<String> EXPECTED_FILE_SET = new HashSet<String>() {
         {
             add("test.xml");
             add("employee.json");
@@ -23,29 +23,49 @@ public class ListFilesUnitTest {
             add("country.txt");
         }
     };
+    private static final int DEPTH = 1;
+    private final ListFiles listFiles = new ListFiles();
 
     @Test
-    public void givenDir_whenUsingJAVAIO_thenListAllFiles() {
-        assertEquals(EXPECTED_FILE_LIST, listFiles.listFilesUsingJavaIO(DIRECTORY));
+    public void givenAValidDirectoryWhenUsingJavaIOThenReturnSetOfFileNames() {
+        assertThat(listFiles.listFilesUsingJavaIO(VALID_DIRECTORY))
+                .isEqualTo(EXPECTED_FILE_SET);
     }
 
     @Test
-    public void givenDir_whenUsingFilesList_thenListAllFiles() throws IOException {
-        assertEquals(EXPECTED_FILE_LIST, listFiles.listFilesUsingFilesList(DIRECTORY));
+    public void givenAInvalidValidDirectoryWhenUsingJavaIOThenThrowsNullPointerExceptino() {
+        assertThrows(NullPointerException.class,
+                () -> listFiles.listFilesUsingJavaIO(INVALID_DIRECTORY));
     }
 
     @Test
-    public void givenDir_whenWalkingTree_thenListAllFiles() throws IOException {
-        assertEquals(EXPECTED_FILE_LIST, listFiles.listFilesUsingFileWalk(DIRECTORY,DEPTH));
+    public void givenAValidDirectoryWhenUsingFilesListThenReturnSetOfFileNames() throws IOException {
+        assertThat(listFiles.listFilesUsingFilesList(VALID_DIRECTORY))
+                .isEqualTo(EXPECTED_FILE_SET);
     }
 
     @Test
-    public void givenDir_whenWalkingTreeWithVisitor_thenListAllFiles() throws IOException {
-        assertEquals(EXPECTED_FILE_LIST, listFiles.listFilesUsingFileWalkAndVisitor(DIRECTORY));
+    public void givenAValidDirectoryWhenUsingFilesWalkWithDepth1ThenReturnSetOfFileNames() throws IOException {
+        assertThat(listFiles.listFilesUsingFileWalk(VALID_DIRECTORY, DEPTH))
+                .isEqualTo(EXPECTED_FILE_SET);
     }
 
     @Test
-    public void givenDir_whenUsingDirectoryStream_thenListAllFiles() throws IOException {
-        assertEquals(EXPECTED_FILE_LIST, listFiles.listFilesUsingDirectoryStream(DIRECTORY));
+    public void givenAValidDirectoryWhenUsingFilesWalkFileTreeThenReturnSetOfFileNames() throws IOException {
+        assertThat(listFiles.listFilesUsingFileWalkAndVisitor(VALID_DIRECTORY))
+                .isEqualTo(EXPECTED_FILE_SET);
+    }
+
+    @Test
+    public void givenAValidDirectoryWhenUsingDirectoryStreamThenReturnSetOfFileNames() throws IOException {
+        assertThat(listFiles.listFilesUsingDirectoryStream(VALID_DIRECTORY))
+                .isEqualTo(EXPECTED_FILE_SET);
+    }
+
+    @Test
+    public void givenAValidFileWhenUsingFilesWalkWithDepth1ThenReturnSetOfTheSameFile() throws IOException {
+        Set<String> expectedFileSet = Collections.singleton("test.xml");
+        String filePathString = "src/test/resources/listFilesUnitTestFolder/test.xml";
+        assertEquals(expectedFileSet, listFiles.listFilesUsingFileWalk(filePathString, DEPTH));
     }
 }