Add snapshot tests for resolver (#5404)

## Summary This PR adds some snapshot tests for the resolver based on executing resolutions within a "mock" of the Airflow repo (that is: a folder that contains a subset of the repo's files, but all empty, and with an only-partially-complete virtual environment). It's intended to act as a lightweight integration test, to enable us to test resolutions on a "real" project without adding a dependency on Airflow itself.
2025-08-28 06:14:27 +00:00 · 2023-06-28 09:38:51 -04:00 · 2023-06-28 09:38:51 -04:00 · 6587fb844a
commit 6587fb844a
parent a68a86e18b
32 changed files with 391 additions and 22 deletions
--- a/crates/ruff_python_resolver/src/implicit_imports.rs
+++ b/crates/ruff_python_resolver/src/implicit_imports.rs
@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::BTreeMap;
 use std::ffi::OsStr;
 use std::fs;
 use std::path::{Path, PathBuf};
@ -24,8 +24,8 @@ pub(crate) struct ImplicitImport {
 }

 /// Find the "implicit" imports within the namespace package at the given path.
-pub(crate) fn find(dir_path: &Path, exclusions: &[&Path]) -> HashMap<String, ImplicitImport> {
-    let mut implicit_imports = HashMap::new();
+pub(crate) fn find(dir_path: &Path, exclusions: &[&Path]) -> BTreeMap<String, ImplicitImport> {
+    let mut implicit_imports = BTreeMap::new();

    // Enumerate all files and directories in the path, expanding links.
    let Ok(entries) = fs::read_dir(dir_path) else {
@ -128,14 +128,14 @@ pub(crate) fn find(dir_path: &Path, exclusions: &[&Path]) -> HashMap<String, Imp

 /// Filter a map of implicit imports to only include those that were actually imported.
 pub(crate) fn filter(
-    implicit_imports: &HashMap<String, ImplicitImport>,
+    implicit_imports: &BTreeMap<String, ImplicitImport>,
    imported_symbols: &[String],
-) -> Option<HashMap<String, ImplicitImport>> {
+) -> Option<BTreeMap<String, ImplicitImport>> {
    if implicit_imports.is_empty() || imported_symbols.is_empty() {
        return None;
    }

-    let mut filtered_imports = HashMap::new();
+    let mut filtered_imports = BTreeMap::new();
    for implicit_import in implicit_imports.values() {
        if imported_symbols.contains(&implicit_import.name) {
            filtered_imports.insert(implicit_import.name.clone(), implicit_import.clone());
--- a/crates/ruff_python_resolver/src/import_result.rs
+++ b/crates/ruff_python_resolver/src/import_result.rs
@ -1,9 +1,9 @@
 //! Interface that describes the output of the import resolver.

-use crate::implicit_imports::ImplicitImport;
-use std::collections::HashMap;
+use std::collections::BTreeMap;
 use std::path::PathBuf;

+use crate::implicit_imports::ImplicitImport;
 use crate::py_typed::PyTypedInfo;

 #[derive(Debug, Clone, PartialEq, Eq)]
@ -69,11 +69,11 @@ pub(crate) struct ImportResult {

    /// A map from file to resolved path, for all implicitly imported
    /// modules that are part of a namespace package.
-    pub(crate) implicit_imports: HashMap<String, ImplicitImport>,
+    pub(crate) implicit_imports: BTreeMap<String, ImplicitImport>,

    /// Any implicit imports whose symbols were explicitly imported (i.e., via
    /// a `from x import y` statement).
-    pub(crate) filtered_implicit_imports: HashMap<String, ImplicitImport>,
+    pub(crate) filtered_implicit_imports: BTreeMap<String, ImplicitImport>,

    /// If the import resolved to a type hint (i.e., a `.pyi` file), then
    /// a non-type-hint resolution will be stored here.
@ -105,8 +105,8 @@ impl ImportResult {
            is_stdlib_typeshed_file: false,
            is_third_party_typeshed_file: false,
            is_local_typings_file: false,
-            implicit_imports: HashMap::default(),
-            filtered_implicit_imports: HashMap::default(),
+            implicit_imports: BTreeMap::default(),
+            filtered_implicit_imports: BTreeMap::default(),
            non_stub_import_result: None,
            py_typed_info: None,
            package_directory: None,
--- a/crates/ruff_python_resolver/src/lib.rs
+++ b/crates/ruff_python_resolver/src/lib.rs
@ -12,5 +12,3 @@ mod python_platform;
 mod python_version;
 mod resolver;
 mod search;
-
-pub(crate) const SITE_PACKAGES: &str = "site-packages";
--- a/crates/ruff_python_resolver/src/resolver.rs
+++ b/crates/ruff_python_resolver/src/resolver.rs
@ -1,6 +1,6 @@
 //! Resolves Python imports to their corresponding files on disk.

-use std::collections::HashMap;
+use std::collections::BTreeMap;
 use std::path::{Path, PathBuf};

 use log::debug;
@ -36,7 +36,7 @@ fn resolve_module_descriptor(
    let mut is_stub_package = false;
    let mut is_stub_file = false;
    let mut is_native_lib = false;
-    let mut implicit_imports = HashMap::new();
+    let mut implicit_imports = BTreeMap::new();
    let mut package_directory = None;
    let mut py_typed_info = None;

@ -194,7 +194,7 @@ fn resolve_module_descriptor(
        is_third_party_typeshed_file: false,
        is_local_typings_file: false,
        implicit_imports,
-        filtered_implicit_imports: HashMap::default(),
+        filtered_implicit_imports: BTreeMap::default(),
        non_stub_import_result: None,
        py_typed_info,
        package_directory,
@ -424,7 +424,7 @@ fn resolve_best_absolute_import<Host: host::Host>(
 /// are all satisfied by submodules (as listed in the implicit imports).
 fn is_namespace_package_resolved(
    module_descriptor: &ImportModuleDescriptor,
-    implicit_imports: &HashMap<String, ImplicitImport>,
+    implicit_imports: &BTreeMap<String, ImplicitImport>,
 ) -> bool {
    if !module_descriptor.imported_symbols.is_empty() {
        // Pyright uses `!Array.from(moduleDescriptor.importedSymbols.keys()).some((symbol) => implicitImports.has(symbol))`.
@ -774,6 +774,7 @@ fn resolve_import<Host: host::Host>(

 #[cfg(test)]
 mod tests {
+    use insta::assert_debug_snapshot;
    use std::fs::{create_dir_all, File};
    use std::io::{self, Write};
    use std::path::{Path, PathBuf};
@ -823,6 +824,8 @@ mod tests {
        library: Option<PathBuf>,
        stub_path: Option<PathBuf>,
        typeshed_path: Option<PathBuf>,
+        venv_path: Option<PathBuf>,
+        venv: Option<PathBuf>,
    }

    fn resolve_options(
@ -836,6 +839,8 @@ mod tests {
            library,
            stub_path,
            typeshed_path,
+            venv_path,
+            venv,
        } = options;

        let execution_environment = ExecutionEnvironment {
@ -860,8 +865,8 @@ mod tests {
        let config = Config {
            typeshed_path,
            stub_path,
-            venv_path: None,
-            venv: None,
+            venv_path,
+            venv,
        };

        let host = host::StaticHost::new(if let Some(library) = library {
@ -1545,4 +1550,109 @@ mod tests {

        Ok(())
    }
+
+    #[test]
+    fn airflow_standard_library() {
+        setup();
+
+        let root = PathBuf::from("./resources/test/airflow");
+        let source_file = root.join("airflow/api/common/mark_tasks.py");
+
+        let result = resolve_options(
+            source_file,
+            "os",
+            root.clone(),
+            ResolverOptions {
+                venv_path: Some(root),
+                venv: Some(PathBuf::from("venv")),
+                ..Default::default()
+            },
+        );
+
+        assert_debug_snapshot!(result);
+    }
+
+    #[test]
+    fn airflow_first_party() {
+        setup();
+
+        let root = PathBuf::from("./resources/test/airflow");
+        let source_file = root.join("airflow/api/common/mark_tasks.py");
+
+        let result = resolve_options(
+            source_file,
+            "airflow.jobs.scheduler_job_runner",
+            root.clone(),
+            ResolverOptions {
+                venv_path: Some(root),
+                venv: Some(PathBuf::from("venv")),
+                ..Default::default()
+            },
+        );
+
+        assert_debug_snapshot!(result);
+    }
+
+    #[test]
+    fn airflow_stub_file() {
+        setup();
+
+        let root = PathBuf::from("./resources/test/airflow");
+        let source_file = root.join("airflow/api/common/mark_tasks.py");
+
+        let result = resolve_options(
+            source_file,
+            "airflow.compat.functools",
+            root.clone(),
+            ResolverOptions {
+                venv_path: Some(root),
+                venv: Some(PathBuf::from("venv")),
+                ..Default::default()
+            },
+        );
+
+        assert_debug_snapshot!(result);
+    }
+
+    #[test]
+    fn airflow_namespace_package() {
+        setup();
+
+        let root = PathBuf::from("./resources/test/airflow");
+        let source_file = root.join("airflow/api/common/mark_tasks.py");
+
+        let result = resolve_options(
+            source_file,
+            "airflow.providers.google.cloud.hooks.gcs",
+            root.clone(),
+            ResolverOptions {
+                venv_path: Some(root),
+                venv: Some(PathBuf::from("venv")),
+                ..Default::default()
+            },
+        );
+
+        assert_debug_snapshot!(result);
+    }
+
+    #[test]
+    fn airflow_third_party() {
+        setup();
+
+        let root = PathBuf::from("./resources/test/airflow");
+        let source_file = root.join("airflow/api/common/mark_tasks.py");
+
+        let result = resolve_options(
+            source_file,
+            "sqlalchemy.orm",
+            root.clone(),
+            ResolverOptions {
+                venv_path: Some(root),
+                venv: Some(PathBuf::from("venv")),
+                ..Default::default()
+            },
+        );
+
+        assert_debug_snapshot!(result);
+    }
 }
--- a/crates/ruff_python_resolver/src/search.rs
+++ b/crates/ruff_python_resolver/src/search.rs
@ -8,9 +8,11 @@ use std::path::{Path, PathBuf};
 use log::debug;

 use crate::config::Config;
+use crate::host;
 use crate::module_descriptor::ImportModuleDescriptor;
 use crate::python_version::PythonVersion;
-use crate::{host, SITE_PACKAGES};
+
+const SITE_PACKAGES: &str = "site-packages";

 /// Find the `site-packages` directory for the specified Python version.
 fn find_site_packages_path(
--- a/crates/ruff_python_resolver/src/snapshots/ruff_python_resolverresolvertests__airflow_first_party.snap
+++ b/crates/ruff_python_resolver/src/snapshots/ruff_python_resolverresolvertests__airflow_first_party.snap
@ -0,0 +1,33 @@
+---
+source: crates/ruff_python_resolver/src/resolver.rs
+expression: result
+---
+ImportResult {
+    is_relative: false,
+    is_import_found: true,
+    is_partly_resolved: false,
+    is_namespace_package: false,
+    is_init_file_present: true,
+    is_stub_package: false,
+    import_type: Local,
+    resolved_paths: [
+        "./resources/test/airflow/airflow/__init__.py",
+        "./resources/test/airflow/airflow/jobs/__init__.py",
+        "./resources/test/airflow/airflow/jobs/scheduler_job_runner.py",
+    ],
+    search_path: Some(
+        "./resources/test/airflow",
+    ),
+    is_stub_file: false,
+    is_native_lib: false,
+    is_stdlib_typeshed_file: false,
+    is_third_party_typeshed_file: false,
+    is_local_typings_file: false,
+    implicit_imports: {},
+    filtered_implicit_imports: {},
+    non_stub_import_result: None,
+    py_typed_info: None,
+    package_directory: Some(
+        "./resources/test/airflow/airflow",
+    ),
+}
--- a/crates/ruff_python_resolver/src/snapshots/ruff_python_resolverresolvertests__airflow_namespace_package.snap
+++ b/crates/ruff_python_resolver/src/snapshots/ruff_python_resolverresolvertests__airflow_namespace_package.snap
@ -0,0 +1,36 @@
+---
+source: crates/ruff_python_resolver/src/resolver.rs
+expression: result
+---
+ImportResult {
+    is_relative: false,
+    is_import_found: true,
+    is_partly_resolved: false,
+    is_namespace_package: true,
+    is_init_file_present: true,
+    is_stub_package: false,
+    import_type: Local,
+    resolved_paths: [
+        "./resources/test/airflow/airflow/__init__.py",
+        "",
+        "./resources/test/airflow/airflow/providers/google/__init__.py",
+        "./resources/test/airflow/airflow/providers/google/cloud/__init__.py",
+        "./resources/test/airflow/airflow/providers/google/cloud/hooks/__init__.py",
+        "./resources/test/airflow/airflow/providers/google/cloud/hooks/gcs.py",
+    ],
+    search_path: Some(
+        "./resources/test/airflow",
+    ),
+    is_stub_file: false,
+    is_native_lib: false,
+    is_stdlib_typeshed_file: false,
+    is_third_party_typeshed_file: false,
+    is_local_typings_file: false,
+    implicit_imports: {},
+    filtered_implicit_imports: {},
+    non_stub_import_result: None,
+    py_typed_info: None,
+    package_directory: Some(
+        "./resources/test/airflow/airflow",
+    ),
+}
--- a/crates/ruff_python_resolver/src/snapshots/ruff_python_resolverresolvertests__airflow_standard_library.snap
+++ b/crates/ruff_python_resolver/src/snapshots/ruff_python_resolverresolvertests__airflow_standard_library.snap
@ -0,0 +1,25 @@
+---
+source: crates/ruff_python_resolver/src/resolver.rs
+expression: result
+---
+ImportResult {
+    is_relative: false,
+    is_import_found: false,
+    is_partly_resolved: false,
+    is_namespace_package: false,
+    is_init_file_present: false,
+    is_stub_package: false,
+    import_type: Local,
+    resolved_paths: [],
+    search_path: None,
+    is_stub_file: false,
+    is_native_lib: false,
+    is_stdlib_typeshed_file: false,
+    is_third_party_typeshed_file: false,
+    is_local_typings_file: false,
+    implicit_imports: {},
+    filtered_implicit_imports: {},
+    non_stub_import_result: None,
+    py_typed_info: None,
+    package_directory: None,
+}
--- a/crates/ruff_python_resolver/src/snapshots/ruff_python_resolverresolvertests__airflow_stub_file.snap
+++ b/crates/ruff_python_resolver/src/snapshots/ruff_python_resolverresolvertests__airflow_stub_file.snap
@ -0,0 +1,63 @@
+---
+source: crates/ruff_python_resolver/src/resolver.rs
+expression: result
+---
+ImportResult {
+    is_relative: false,
+    is_import_found: true,
+    is_partly_resolved: false,
+    is_namespace_package: false,
+    is_init_file_present: true,
+    is_stub_package: false,
+    import_type: Local,
+    resolved_paths: [
+        "./resources/test/airflow/airflow/__init__.py",
+        "./resources/test/airflow/airflow/compat/__init__.py",
+        "./resources/test/airflow/airflow/compat/functools.pyi",
+    ],
+    search_path: Some(
+        "./resources/test/airflow",
+    ),
+    is_stub_file: true,
+    is_native_lib: false,
+    is_stdlib_typeshed_file: false,
+    is_third_party_typeshed_file: false,
+    is_local_typings_file: false,
+    implicit_imports: {},
+    filtered_implicit_imports: {},
+    non_stub_import_result: Some(
+        ImportResult {
+            is_relative: false,
+            is_import_found: true,
+            is_partly_resolved: false,
+            is_namespace_package: false,
+            is_init_file_present: true,
+            is_stub_package: false,
+            import_type: Local,
+            resolved_paths: [
+                "./resources/test/airflow/airflow/__init__.py",
+                "./resources/test/airflow/airflow/compat/__init__.py",
+                "./resources/test/airflow/airflow/compat/functools.py",
+            ],
+            search_path: Some(
+                "./resources/test/airflow",
+            ),
+            is_stub_file: false,
+            is_native_lib: false,
+            is_stdlib_typeshed_file: false,
+            is_third_party_typeshed_file: false,
+            is_local_typings_file: false,
+            implicit_imports: {},
+            filtered_implicit_imports: {},
+            non_stub_import_result: None,
+            py_typed_info: None,
+            package_directory: Some(
+                "./resources/test/airflow/airflow",
+            ),
+        },
+    ),
+    py_typed_info: None,
+    package_directory: Some(
+        "./resources/test/airflow/airflow",
+    ),
+}
--- a/crates/ruff_python_resolver/src/snapshots/ruff_python_resolverresolvertests__airflow_third_party.snap
+++ b/crates/ruff_python_resolver/src/snapshots/ruff_python_resolverresolvertests__airflow_third_party.snap
@ -0,0 +1,54 @@
+---
+source: crates/ruff_python_resolver/src/resolver.rs
+expression: result
+---
+ImportResult {
+    is_relative: false,
+    is_import_found: true,
+    is_partly_resolved: false,
+    is_namespace_package: false,
+    is_init_file_present: true,
+    is_stub_package: false,
+    import_type: ThirdParty,
+    resolved_paths: [
+        "./resources/test/airflow/venv/Lib/python3.11/site-packages/sqlalchemy/__init__.py",
+        "./resources/test/airflow/venv/Lib/python3.11/site-packages/sqlalchemy/orm/__init__.py",
+    ],
+    search_path: Some(
+        "./resources/test/airflow/venv/Lib/python3.11/site-packages",
+    ),
+    is_stub_file: false,
+    is_native_lib: false,
+    is_stdlib_typeshed_file: false,
+    is_third_party_typeshed_file: false,
+    is_local_typings_file: false,
+    implicit_imports: {
+        "base": ImplicitImport {
+            is_stub_file: false,
+            is_native_lib: false,
+            name: "base",
+            path: "./resources/test/airflow/venv/Lib/python3.11/site-packages/sqlalchemy/orm/base.py",
+            py_typed: None,
+        },
+        "dependency": ImplicitImport {
+            is_stub_file: false,
+            is_native_lib: false,
+            name: "dependency",
+            path: "./resources/test/airflow/venv/Lib/python3.11/site-packages/sqlalchemy/orm/dependency.py",
+            py_typed: None,
+        },
+        "query": ImplicitImport {
+            is_stub_file: false,
+            is_native_lib: false,
+            name: "query",
+            path: "./resources/test/airflow/venv/Lib/python3.11/site-packages/sqlalchemy/orm/query.py",
+            py_typed: None,
+        },
+    },
+    filtered_implicit_imports: {},
+    non_stub_import_result: None,
+    py_typed_info: None,
+    package_directory: Some(
+        "./resources/test/airflow/venv/Lib/python3.11/site-packages/sqlalchemy",
+    ),
+}