diff --git a/Makefile b/Makefile
index beae4f9..c29a4c8 100644
--- a/Makefile
+++ b/Makefile
@@ -23,3 +23,6 @@ run_test_parsing:
 
 run_test_dfs:
 	PYTHONPATH=src uv run pytest tests/test_Depth.py
+
+run_test_maze_gen:
+	PYTHONPATH=src uv run pytest tests/test_MazeGenerator.py
diff --git a/src/amaz_lib/MazeGenerator.py b/src/amaz_lib/MazeGenerator.py
index 4240759..daeede7 100644
--- a/src/amaz_lib/MazeGenerator.py
+++ b/src/amaz_lib/MazeGenerator.py
@@ -94,26 +94,22 @@ class DepthFirstSearch:
         path = list()
         w_h = (width, height)
         coord = (0, 0)
+        x, y = coord
+
         while len(visited) < width * height:
-            print(f"visited {len(visited)}")
-            x, y = coord
-            rand_steps = DepthFirstSearch.random_cells(visited, coord, w_h)
-            if len(rand_steps) == 0:
-                path = DepthFirstSearch.back_on_step(path, w_h, visited)
-                coord = path[-1]
-                rand_steps = DepthFirstSearch.random_cells(path, coord, w_h)
-                print(f"coord {coord}")
-                print(f"visited = {visited}")
-                print(f" rand steps {rand_steps}")
-                print(f"path = {path}")
-                x, y = coord
-            wall = DepthFirstSearch.next_step(rand_steps)
-            wall_r = DepthFirstSearch.reverse_path(wall)
-            maze[y][x] = DepthFirstSearch.broken_wall(maze[y][x], wall)
             visited = DepthFirstSearch.add_cell_visited(coord, visited)
             path = DepthFirstSearch.add_cell_visited(coord, path)
+            random_c = DepthFirstSearch.random_cells(visited, coord, w_h)
+            if len(random_c) == 0:
+                path = DepthFirstSearch.back_on_step(path, w_h, visited)
+                if path:
+                    coord = path[-1]
+                random_c = DepthFirstSearch.random_cells(path, coord, w_h)
+                x, y = coord
+            wall = DepthFirstSearch.next_step(random_c)
+            maze[y][x] = DepthFirstSearch.broken_wall(maze[y][x], wall)
             coord = DepthFirstSearch.next_cell(x, y, wall)
-            print(f"coord 2 {coord}")
+            wall_r = DepthFirstSearch.reverse_path(wall)
             x, y = coord
             maze[y][x] = DepthFirstSearch.broken_wall(maze[y][x], wall_r)
         return maze
@@ -196,11 +192,10 @@ class DepthFirstSearch:
         last = path[-1]
         r_cells = DepthFirstSearch.random_cells(visited,  last, w_h)
         while len(path) > 0:
-            print(f"path {len(path)}")
             path.pop()
-            last = path[-1]
+            if path:
+                last = path[-1]
             r_cells = DepthFirstSearch.random_cells(visited,  last, w_h)
             if r_cells:
-                print(f"cells {len(r_cells)}")
                 break
         return path
diff --git a/tests/test_MazeGenerator.py b/tests/test_MazeGenerator.py
index cb0aa6b..a9bf22b 100644
--- a/tests/test_MazeGenerator.py
+++ b/tests/test_MazeGenerator.py
@@ -1,11 +1,18 @@
 import numpy
-from amaz_lib.MazeGenerator import Kruskal
+from amaz_lib.MazeGenerator import DepthFirstSearch
 
 
-def test_kruskal_output_shape() -> None:
-    generator = Kruskal()
-    maze = numpy.array([])
-    for output in generator.generator(10, 10):
-        maze = output
+class TestMazeGenerator:
 
-    assert maze.shape == (10, 10)
+    # def test_kruskal_output_shape() -> None:
+    #     generator = Kruskal()
+    #     maze = numpy.array([])
+    #     for output in generator.generator(10, 10):
+    #         maze = output
+
+    #     assert maze.shape == (10, 10)
+
+    def test_generator(self):
+        maze = numpy.array([])
+        maze = DepthFirstSearch.generator(10, 10)
+        assert maze.shape == (10, 10)