diff --git a/tests/test_cli.py b/tests/test_cli.py
index d55d1e6a..3bf00ffb 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -7,7 +7,7 @@
 from io import StringIO
 from unittest import mock
 
-from chipflow_lib import ChipFlowError
+from chipflow import ChipFlowError
 from chipflow.cli import run
 from chipflow.config_models import Config, ChipFlowConfig
 
diff --git a/tests/test_init.py b/tests/test_init.py
index 3ec38093..f3cfae3d 100644
--- a/tests/test_init.py
+++ b/tests/test_init.py
@@ -7,12 +7,11 @@
 from pathlib import Path
 from unittest import mock
 
-from chipflow_lib import (
-    ChipFlowError,
+from chipflow import ChipFlowError, ensure_chipflow_root
+from chipflow.cli import (
     _get_cls_by_reference,
-    _ensure_chipflow_root,
-    _parse_config
-)
+    _parse_config)
+
 from chipflow.config.parser import _parse_config_file
 from chipflow.config_models import Config, ChipFlowConfig
 # Process is not part of the public API, so we won't test it here
@@ -75,8 +74,8 @@ def test_ensure_chipflow_root_already_set(self):
         os.environ["CHIPFLOW_ROOT"] = "/test/path"
         sys.path = ["/some/other/path"]
 
-        _ensure_chipflow_root.root = None  #type: ignore
-        result = _ensure_chipflow_root()
+        ensure_chipflow_root.root = None  #type: ignore
+        result = ensure_chipflow_root()
 
         self.assertEqual(result, Path("/test/path"))
         self.assertIn("/test/path", sys.path)
@@ -85,10 +84,10 @@ def test_ensure_chipflow_root_not_set(self):
         """Test _ensure_chipflow_root when CHIPFLOW_ROOT is not set"""
         if "CHIPFLOW_ROOT" in os.environ:
             del os.environ["CHIPFLOW_ROOT"]
-        _ensure_chipflow_root.root = None  #type: ignore
+        ensure_chipflow_root.root = None  #type: ignore
 
         with mock.patch("os.getcwd", return_value="/mock/cwd"):
-            result = _ensure_chipflow_root()
+            result = ensure_chipflow_root()
 
             self.assertEqual(result, Path("/mock/cwd"))
             self.assertEqual(os.environ["CHIPFLOW_ROOT"], "/mock/cwd")