diff --git a/src/mazegen/MazeGenerator.py b/src/mazegen/MazeGenerator.py
index 866763a..a04c568 100644
--- a/src/mazegen/MazeGenerator.py
+++ b/src/mazegen/MazeGenerator.py
@@ -2,7 +2,7 @@ from abc import ABC, abstractmethod
 from typing import Generator, Any
 import numpy as np
 from numpy.typing import NDArray
-from .Cell import Cell
+from mazegen import Cell
 import math
 import random
 
diff --git a/tests/test_Cell.py b/tests/test_Cell.py
index 7b0561a..9024ba4 100644
--- a/tests/test_Cell.py
+++ b/tests/test_Cell.py
@@ -1,4 +1,4 @@
-from amaz_lib.Cell import Cell
+from mazegen import Cell
 
 
 def test_cell_setter_getter() -> None:
diff --git a/tests/test_Depth.py b/tests/test_Depth.py
index 32a6dad..2ae1ed8 100644
--- a/tests/test_Depth.py
+++ b/tests/test_Depth.py
@@ -1,5 +1,5 @@
-from amaz_lib.MazeGenerator import DepthFirstSearch
-from amaz_lib.Cell import Cell
+from mazegen import DepthFirstSearch
+from mazegen import Cell
 import numpy as np
 
 
diff --git a/tests/test_Maze.py b/tests/test_Maze.py
index 282d4b5..a4cb962 100644
--- a/tests/test_Maze.py
+++ b/tests/test_Maze.py
@@ -1,6 +1,6 @@
 import numpy
-from amaz_lib.Cell import Cell
-from amaz_lib.Maze import Maze
+from mazegen import Cell
+from mazegen import Maze
 
 
 def test_maze_setter_getter() -> None:
diff --git a/tests/test_MazeGenerator.py b/tests/test_MazeGenerator.py
index 9198fe9..5c6796b 100644
--- a/tests/test_MazeGenerator.py
+++ b/tests/test_MazeGenerator.py
@@ -1,5 +1,5 @@
 import numpy
-from amaz_lib.MazeGenerator import DepthFirstSearch
+from mazegen import DepthFirstSearch
 
 
 class TestMazeGenerator:
diff --git a/tests/test_MazeSolver.py b/tests/test_MazeSolver.py
index 4ec0061..dd14e3d 100644
--- a/tests/test_MazeSolver.py
+++ b/tests/test_MazeSolver.py
@@ -1,6 +1,6 @@
-from amaz_lib.Cell import Cell
+from mazegen import Cell
 import numpy as np
-from amaz_lib import AStar, Maze
+from mazegen import AStar, Maze
 
 
 def test_solver() -> None: