Move tests to the matching modules

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))

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_notebooks_dbutils_notebook_run.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))

tests/unit/source_code/test_graph.py

@@ -1,4 +1,3 @@
-from collections.abc import Iterable
 from pathlib import Path
 
 import pytest
@@ -11,10 +10,7 @@
     DependencyProblem,
     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:
@@ -195,31 +191,6 @@ def test_graph_builds_inherited_context(mock_path_lookup, simple_dependency_reso
     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)
-
-
 def test_dependency_problem_has_path_missing_by_default() -> None:
     problem = DependencyProblem("code", "message")
     assert problem.has_missing_path()

tests/unit/source_code/test_jobs.py

@@ -20,15 +20,13 @@
 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,
     DependencyGraph,
     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_notebooks_dbutils_notebook_run.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)