Replace MockInstallation with MockPathLookup for testing fixtures

tests/unit/__init__.py

@@ -22,6 +22,7 @@
 from databricks.labs.ucx.source_code.graph import SourceContainer
 
 logging.getLogger("tests").setLevel("DEBUG")
+logger = logging.getLogger(__name__)
 
 DEFAULT_CONFIG = {
     "config.yml": {
@@ -121,16 +122,6 @@ def _id_list(cls: type, ids=None):
     return [installation.load(cls, filename=_) for _ in ids]
 
 
-def _load_sources(cls: type, *filenames: str):
-    # TODO: remove the usage of it in favor of MockPathLookup
-    if not filenames:
-        return []
-    installation = MockInstallation(DEFAULT_CONFIG | {_: _load_source(f'{_FOLDERS[cls]}/{_}') for _ in filenames})
-    # cleanly avoid mypy error
-    setattr(installation, "_unmarshal_type", lambda as_dict, filename, type_ref: as_dict)
-    return [installation.load(cls, filename=_) for _ in filenames]
-
-
 def _samples_path(cls: type) -> str:
     return (__dir / _FOLDERS[cls]).as_posix()
 

tests/unit/conftest.py

@@ -219,3 +219,16 @@ def simple_dependency_resolver(mock_path_lookup: PathLookup) -> DependencyResolv
     notebook_resolver = NotebookResolver(NotebookLoader())
     import_resolver = ImportFileResolver(FileLoader(), allow_list)
     return DependencyResolver(library_resolver, notebook_resolver, import_resolver, import_resolver, mock_path_lookup)
+
+
+def _load_sources(*filenames: str):
+    # Load the contents of the files into a list of strings
+    contents = []
+    path_lookup = MockPathLookup()
+    for filename in filenames:
+        path = path_lookup.resolve(Path(filename))
+        # Read local file into a string
+        if path:
+            with open(path.as_posix(), 'r', encoding='utf-8') as f:
+                contents.append(f.read())
+    return contents

tests/unit/source_code/test_notebook.py

@@ -5,7 +5,7 @@
 from databricks.sdk.service.workspace import Language, ObjectType, ObjectInfo
 
 from databricks.labs.ucx.source_code.base import CurrentSessionState
-from databricks.labs.ucx.source_code.graph import DependencyGraph, SourceContainer, DependencyResolver
+from databricks.labs.ucx.source_code.graph import DependencyGraph, DependencyResolver
 from databricks.labs.ucx.source_code.known import KnownList
 from databricks.labs.ucx.source_code.linters.files import ImportFileResolver, FileLoader
 from databricks.labs.ucx.source_code.linters.imports import DbutilsPyLinter
@@ -16,7 +16,7 @@
     NotebookLoader,
 )
 from databricks.labs.ucx.source_code.python_libraries import PythonLibraryResolver
-from tests.unit import _load_sources
+from tests.unit.conftest import _load_sources
 
 # fmt: off
 # the following samples are real samples from https://github.com/databricks-industry-solutions
@@ -78,7 +78,7 @@
 )
 def test_notebook_splits_source_into_cells(source: tuple[str, Language, list[str]]) -> None:
     path = source[0]
-    sources: list[str] = _load_sources(SourceContainer, path)
+    sources: list[str] = _load_sources(path)
     assert len(sources) == 1
     notebook = Notebook.parse(Path(path), sources[0], source[1])
     assert notebook is not None
@@ -98,7 +98,7 @@ def test_notebook_splits_source_into_cells(source: tuple[str, Language, list[str
 )
 def test_notebook_rebuilds_same_code(source: tuple[str, Language, list[str]]) -> None:
     path = source[0]
-    sources: list[str] = _load_sources(SourceContainer, path)
+    sources: list[str] = _load_sources(path)
     assert len(sources) == 1
     notebook = Notebook.parse(Path(path), sources[0], source[1])
     assert notebook is not None
@@ -121,7 +121,7 @@ def test_notebook_rebuilds_same_code(source: tuple[str, Language, list[str]]) ->
 )
 def test_notebook_generates_runnable_cells(source: tuple[str, Language, list[str]]) -> None:
     path = source[0]
-    sources: list[str] = _load_sources(SourceContainer, path)
+    sources: list[str] = _load_sources(path)
     assert len(sources) == 1
     notebook = Notebook.parse(Path(path), sources[0], source[1])
     assert notebook is not None