diff --git a/tests/unit/source_code/linters/test_directfs.py b/tests/unit/source_code/linters/test_directfs.py index d401280372..02790ad346 100644 --- a/tests/unit/source_code/linters/test_directfs.py +++ b/tests/unit/source_code/linters/test_directfs.py @@ -1,18 +1,11 @@ -from collections.abc import Iterable -from pathlib import Path -from unittest.mock import create_autospec - import pytest -from databricks.labs.ucx.source_code.base import Deprecation, Advice, CurrentSessionState, Failure, DirectFsAccess -from databricks.labs.ucx.source_code.graph import DependencyGraph -from databricks.labs.ucx.source_code.linters.graph_walkers import DfsaCollectorWalker +from databricks.labs.ucx.source_code.base import Deprecation, Advice, CurrentSessionState, Failure from databricks.labs.ucx.source_code.linters.directfs import ( DIRECT_FS_ACCESS_PATTERNS, DirectFsAccessPyLinter, DirectFsAccessSqlLinter, ) -from databricks.labs.ucx.source_code.notebooks.cells import CellLanguage @pytest.mark.parametrize( @@ -152,18 +145,3 @@ def test_dfsa_queries_failure(query: str) -> None: end_col=1024, ), ] - - -class _TestCollectorWalker(DfsaCollectorWalker): - # inherit from DfsaCollectorWalker because it's public - - def collect_from_source(self, language: CellLanguage) -> Iterable[DirectFsAccess]: - return self._collect_from_source("empty", language, Path(""), None) - - -@pytest.mark.parametrize("language", list(iter(CellLanguage))) -def test_collector_supports_all_cell_languages(language, mock_path_lookup, migration_index): - graph = create_autospec(DependencyGraph) - graph.assert_not_called() - collector = _TestCollectorWalker(graph, set(), mock_path_lookup, CurrentSessionState(), migration_index) - list(collector.collect_from_source(language)) diff --git a/tests/unit/source_code/linters/test_graph_walkers.py b/tests/unit/source_code/linters/test_graph_walkers.py new file mode 100644 index 0000000000..01ad02750c --- /dev/null +++ b/tests/unit/source_code/linters/test_graph_walkers.py @@ -0,0 +1,73 @@ +from collections.abc import Iterable +from pathlib import Path +from unittest.mock import create_autospec + +import pytest + +from databricks.labs.ucx.source_code.base import CurrentSessionState, DirectFsAccess +from databricks.labs.ucx.source_code.graph import Dependency, DependencyGraph +from databricks.labs.ucx.source_code.linters.context import LinterContext +from databricks.labs.ucx.source_code.linters.graph_walkers import ( + DependencyGraphWalker, + LintingWalker, + DfsaCollectorWalker, +) +from databricks.labs.ucx.source_code.notebooks.cells import CellLanguage +from databricks.labs.ucx.source_code.notebooks.loaders import NotebookLoader +from databricks.labs.ucx.source_code.path_lookup import PathLookup +from databricks.labs.ucx.source_code.python.python_ast import Tree + + +def test_graph_walker_captures_lineage(mock_path_lookup, simple_dependency_resolver) -> None: + grand_parent = mock_path_lookup.cwd / "functional/grand_parent_that_magic_runs_parent_that_magic_runs_child.py" + child = mock_path_lookup.cwd / "functional/_child_that_uses_value_from_parent.py" + root_dependency = Dependency(NotebookLoader(), grand_parent) + root_graph = DependencyGraph( + root_dependency, None, simple_dependency_resolver, mock_path_lookup, CurrentSessionState() + ) + container = root_dependency.load(mock_path_lookup) + assert container is not None + container.build_dependency_graph(root_graph) + + class _TestWalker(DependencyGraphWalker): + def _process_dependency( + self, dependency: Dependency, path_lookup: PathLookup, inherited_tree: Tree | None + ) -> Iterable[None]: + if dependency.path.as_posix().endswith(grand_parent.as_posix()): + assert len(self._lineage) == 1 + if dependency.path.as_posix().endswith(child.as_posix()): + assert len(self._lineage) == 3 # there's a parent between grand_parent and child + return [] + + walker = _TestWalker(root_graph, set(), mock_path_lookup) + _ = list(_ for _ in walker) + + +def test_linting_walker_populates_paths(simple_dependency_resolver, mock_path_lookup, migration_index) -> None: + path = mock_path_lookup.resolve(Path("functional/values_across_cells.py")) + root = Dependency(NotebookLoader(), path) + xgraph = DependencyGraph(root, None, simple_dependency_resolver, mock_path_lookup, CurrentSessionState()) + current_session = CurrentSessionState() + walker = LintingWalker( + xgraph, set(), mock_path_lookup, "key", current_session, lambda: LinterContext(migration_index, current_session) + ) + advices = 0 + for advice in walker: + advices += 1 + assert "UNKNOWN" not in advice.path.as_posix() + assert advices + + +class _TestCollectorWalker(DfsaCollectorWalker): + # inherit from DfsaCollectorWalker because it's public + + def collect_from_source(self, language: CellLanguage) -> Iterable[DirectFsAccess]: + return self._collect_from_source("empty", language, Path(""), None) + + +@pytest.mark.parametrize("language", list(iter(CellLanguage))) +def test_collector_supports_all_cell_languages(language, mock_path_lookup, migration_index): + graph = create_autospec(DependencyGraph) + graph.assert_not_called() + collector = _TestCollectorWalker(graph, set(), mock_path_lookup, CurrentSessionState(), migration_index) + list(collector.collect_from_source(language)) diff --git a/tests/unit/source_code/test_graph.py b/tests/unit/source_code/test_graph.py index 47b16f19e3..757f984204 100644 --- a/tests/unit/source_code/test_graph.py +++ b/tests/unit/source_code/test_graph.py @@ -1,4 +1,3 @@ -from collections.abc import Iterable from pathlib import Path import pytest @@ -10,10 +9,7 @@ DependencyGraph, InheritedContext, ) -from databricks.labs.ucx.source_code.linters.graph_walkers import DependencyGraphWalker from databricks.labs.ucx.source_code.notebooks.loaders import NotebookLoader -from databricks.labs.ucx.source_code.path_lookup import PathLookup -from databricks.labs.ucx.source_code.python.python_ast import Tree def test_dependency_graph_registers_library_from_egg(mock_path_lookup, simple_dependency_resolver) -> None: @@ -192,28 +188,3 @@ def test_graph_builds_inherited_context(mock_path_lookup, simple_dependency_reso assert inference_context.tree is not None assert inference_context.tree.has_global("some_table_name") assert not inference_context.tree.has_global("other_table_name") - - -def test_graph_walker_captures_lineage(mock_path_lookup, simple_dependency_resolver) -> None: - grand_parent = mock_path_lookup.cwd / "functional/grand_parent_that_magic_runs_parent_that_magic_runs_child.py" - child = mock_path_lookup.cwd / "functional/_child_that_uses_value_from_parent.py" - root_dependency = Dependency(NotebookLoader(), grand_parent) - root_graph = DependencyGraph( - root_dependency, None, simple_dependency_resolver, mock_path_lookup, CurrentSessionState() - ) - container = root_dependency.load(mock_path_lookup) - assert container is not None - container.build_dependency_graph(root_graph) - - class _TestWalker(DependencyGraphWalker): - def _process_dependency( - self, dependency: Dependency, path_lookup: PathLookup, inherited_tree: Tree | None - ) -> Iterable[None]: - if dependency.path.as_posix().endswith(grand_parent.as_posix()): - assert len(self._lineage) == 1 - if dependency.path.as_posix().endswith(child.as_posix()): - assert len(self._lineage) == 3 # there's a parent between grand_parent and child - return [] - - walker = _TestWalker(root_graph, set(), mock_path_lookup) - _ = list(_ for _ in walker) diff --git a/tests/unit/source_code/test_jobs.py b/tests/unit/source_code/test_jobs.py index 3549a40e22..e5b3866ee7 100644 --- a/tests/unit/source_code/test_jobs.py +++ b/tests/unit/source_code/test_jobs.py @@ -20,7 +20,6 @@ from databricks.sdk.service import compute, jobs, pipelines from databricks.sdk.service.workspace import ExportFormat, ObjectInfo, Language -from databricks.labs.ucx.source_code.linters.context import LinterContext from databricks.labs.ucx.source_code.linters.files import FileLoader, ImportFileResolver from databricks.labs.ucx.source_code.graph import ( Dependency, @@ -28,7 +27,6 @@ DependencyResolver, ) from databricks.labs.ucx.source_code.jobs import JobProblem, WorkflowLinter, WorkflowTaskContainer -from databricks.labs.ucx.source_code.linters.graph_walkers import LintingWalker from databricks.labs.ucx.source_code.notebooks.loaders import NotebookResolver, NotebookLoader from databricks.labs.ucx.source_code.used_table import UsedTablesCrawler @@ -517,21 +515,6 @@ def test_xxx(graph) -> None: ws.assert_not_called() -def test_linting_walker_populates_paths(dependency_resolver, mock_path_lookup, migration_index) -> None: - path = mock_path_lookup.resolve(Path("functional/values_across_cells.py")) - root = Dependency(NotebookLoader(), path) - xgraph = DependencyGraph(root, None, dependency_resolver, mock_path_lookup, CurrentSessionState()) - current_session = CurrentSessionState() - walker = LintingWalker( - xgraph, set(), mock_path_lookup, "key", current_session, lambda: LinterContext(migration_index, current_session) - ) - advices = 0 - for advice in walker: - advices += 1 - assert "UNKNOWN" not in advice.path.as_posix() - assert advices - - def test_workflow_linter_refresh_report(dependency_resolver, mock_path_lookup, migration_index) -> None: ws = create_autospec(WorkspaceClient) ws.workspace.get_status.return_value = ObjectInfo(object_id=123, language=Language.PYTHON)