Add tests for find_sources, gen-source-db, run-tree (#7)

Summary:
Most of the tests test the shared find_sources code directly, but also
add a smoke test each for run-tree and gen-source-db that verify they
are both supporting site-packages.

Pull Request resolved: #7

Reviewed By: martindemello

Differential Revision: D102024363

Pulled By: brittanyrey

fbshipit-source-id: dfa05d6fd773a3ece6d7a76d6698dc08a099cced

Author: msullivan

Cargo.lock

@@ -35,6 +35,7 @@ serde = { version = "1.0.219", features = ["derive", "rc"] }
 serde_json = { version = "1.0.140", features = ["alloc", "float_roundtrip", "raw_value", "unbounded_depth"] }
 starlark_map = "0.13.0"
 static_interner = "0.1.1"
+tempfile = "3.27.0"
 toml = { version = "0.9.12", features = ["preserve_order"] }
 tracing = { version = "0.1.41", features = ["attributes", "valuable"] }
 tracing-subscriber = { version = "0.3.23", features = ["chrono", "env-filter", "json", "local-time", "parking_lot", "registry"] }

src/Cargo.toml

@@ -7,6 +7,7 @@
 
 use std::collections::HashMap;
 use std::collections::HashSet;
+use std::fs;
 use std::process::Command;
 
 use ahash::AHashMap;
@@ -15,6 +16,7 @@ use itertools::Itertools;
 use pyrefly_python::module::Module;
 use pyrefly_python::module_name::ModuleName;
 use rayon::prelude::*;
+use tempfile::TempDir;
 
 use crate::analyzer::analyze;
 use crate::effects::Effect;
@@ -445,3 +447,18 @@ pub fn check_buck_availability() -> bool {
         }
     }
 }
+
+/// Create a new temp directory and write each `(rel_path, contents)` pair
+/// into it, creating intermediate directories as needed. The returned
+/// [`TempDir`] owns the path and deletes it on drop.
+pub fn populate_temp_dir(files: &[(&str, &str)]) -> TempDir {
+    let tmp = TempDir::new().expect("create temp dir");
+    for (rel, contents) in files {
+        let path = tmp.path().join(rel);
+        if let Some(parent) = path.parent() {
+            fs::create_dir_all(parent).expect("create parent dirs");
+        }
+        fs::write(&path, contents).expect("write file");
+    }
+    tmp
+}

src/test_lib.rs

@@ -1,4 +1,4 @@
-# @generated by autocargo from //safer_lazy_imports/lifeguard/tests:[assign,attrs,basic,blocks,calls,constructors,decorators,effects,encoding,exceptions,imports,nested_functions,nested_imports,output,param_mutation,port_test_analyzer,port_test_catch_implicit_imports,port_test_catch_import_cycles,port_test_finalization,port_test_interpreter,port_test_ownership,source_analyzer,star_imports,stubs]
+# @generated by autocargo from //safer_lazy_imports/lifeguard/tests:[assign,attrs,basic,blocks,calls,constructors,decorators,effects,encoding,exceptions,find_sources,gen_source_db,imports,nested_functions,nested_imports,output,param_mutation,port_test_analyzer,port_test_catch_implicit_imports,port_test_catch_import_cycles,port_test_finalization,port_test_interpreter,port_test_ownership,run_tree,source_analyzer,star_imports,stubs]
 
 [package]
 name = "safer_lazy_imports_lifeguard_tests"
@@ -48,6 +48,14 @@ path = "encoding.rs"
 name = "exceptions"
 path = "exceptions.rs"
 
+[[test]]
+name = "find_sources"
+path = "find_sources.rs"
+
+[[test]]
+name = "gen_source_db"
+path = "gen_source_db.rs"
+
 [[test]]
 name = "imports"
 path = "imports.rs"
@@ -92,6 +100,10 @@ path = "port_test_interpreter.rs"
 name = "port_test_ownership"
 path = "port_test_ownership.rs"
 
+[[test]]
+name = "run_tree"
+path = "run_tree.rs"
+
 [[test]]
 name = "source_analyzer"
 path = "source_analyzer.rs"
@@ -106,6 +118,7 @@ path = "stubs.rs"
 
 [dev-dependencies]
 anyhow = "1.0.102"
+clap = { version = "4.6.0", features = ["derive", "env", "string", "unicode", "wrap_help"] }
 lifeguard = { path = "../src" }
 ruff_text_size = { git = "https://github.com/astral-sh/ruff/", rev = "474b00568ad78f02ad8e19b8166cbeb6d69f8511" }
 serde_json = { version = "1.0.140", features = ["alloc", "float_roundtrip", "raw_value", "unbounded_depth"] }

tests/Cargo.toml

@@ -0,0 +1,65 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#[cfg(test)]
+mod tests {
+    use std::collections::BTreeMap;
+    use std::collections::BTreeSet;
+
+    use lifeguard::find_sources::build_source_db;
+    use lifeguard::test_lib::populate_temp_dir;
+
+    fn keys(build_map: &BTreeMap<String, String>) -> BTreeSet<&str> {
+        build_map.keys().map(|s| s.as_str()).collect()
+    }
+
+    #[test]
+    fn test_seeds_py_files_from_input_dir() {
+        let tmp = populate_temp_dir(&[("a.py", ""), ("pkg/__init__.py", ""), ("pkg/b.py", "")]);
+
+        let (build_map, seed_count) = build_source_db(tmp.path(), None).unwrap();
+        assert_eq!(seed_count, 3, "all three .py files are seeded");
+        assert_eq!(
+            keys(&build_map),
+            BTreeSet::from(["a.py", "pkg/__init__.py", "pkg/b.py"]),
+        );
+    }
+
+    #[test]
+    fn test_skips_non_identifier_names() {
+        let tmp = populate_temp_dir(&[
+            ("good.py", ""),
+            // Dir whose name is not a valid identifier: skipped wholesale.
+            (".venv/bad.py", ""),
+            // File whose stem is not a valid identifier: skipped.
+            ("2024-migration.py", ""),
+        ]);
+
+        let (build_map, _) = build_source_db(tmp.path(), None).unwrap();
+        assert_eq!(keys(&build_map), BTreeSet::from(["good.py"]));
+    }
+
+    #[test]
+    fn test_follows_imports_into_site_packages() {
+        let tmp = populate_temp_dir(&[
+            ("proj/main.py", "import foo\n"),
+            ("sp/foo/__init__.py", ""),
+            // Unreachable from main.py's imports — must not be pulled in.
+            ("sp/foo/helper.py", ""),
+            ("sp/unused/__init__.py", ""),
+        ]);
+        let proj = tmp.path().join("proj");
+        let sp = tmp.path().join("sp");
+
+        let (build_map, seed_count) = build_source_db(&proj, Some(&sp)).unwrap();
+        assert_eq!(seed_count, 1, "only main.py is seeded from the project");
+        assert_eq!(
+            keys(&build_map),
+            BTreeSet::from(["main.py", "foo/__init__.py"]),
+        );
+    }
+}

tests/find_sources.rs

@@ -0,0 +1,50 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#[cfg(test)]
+mod tests {
+    use std::collections::BTreeSet;
+    use std::fs;
+
+    use clap::Parser;
+    use lifeguard::commands::gen_source_db::GenSourceDbArgs;
+    use lifeguard::commands::gen_source_db::run;
+    use lifeguard::test_lib::populate_temp_dir;
+    use serde_json::Value;
+
+    #[test]
+    fn test_gen_source_db_writes_build_map_json() {
+        let tmp = populate_temp_dir(&[
+            // Site-packages location is discovered via pyproject.toml rather
+            // than passed on the CLI.
+            (
+                "proj/pyproject.toml",
+                "[lifeguard]\nsite_packages = \"../sp\"\n",
+            ),
+            ("proj/m.py", "import other\n"),
+            ("sp/other/__init__.py", ""),
+        ]);
+        let proj = tmp.path().join("proj");
+        let output = tmp.path().join("db.json");
+
+        let args = GenSourceDbArgs::try_parse_from([
+            "gen-source-db",
+            proj.to_str().unwrap(),
+            output.to_str().unwrap(),
+        ])
+        .unwrap();
+        run(args).unwrap();
+
+        let content = fs::read_to_string(&output).unwrap();
+        let value: Value = serde_json::from_str(&content).unwrap();
+        let build_map = value["build_map"]
+            .as_object()
+            .expect("build_map object in output JSON");
+        let keys: BTreeSet<&str> = build_map.keys().map(|s| s.as_str()).collect();
+        assert_eq!(keys, BTreeSet::from(["m.py", "other/__init__.py"]));
+    }
+}

tests/gen_source_db.rs

@@ -0,0 +1,53 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#[cfg(test)]
+mod tests {
+    use std::collections::BTreeSet;
+    use std::fs;
+
+    use clap::Parser;
+    use lifeguard::commands::run_tree::RunTreeArgs;
+    use lifeguard::commands::run_tree::run;
+    use lifeguard::test_lib::populate_temp_dir;
+    use serde_json::Value;
+
+    #[test]
+    fn test_run_tree_resolves_cli_site_packages() {
+        let tmp = populate_temp_dir(&[
+            ("proj/main.py", "import foo\n"),
+            ("sp/foo/__init__.py", ""),
+            ("sp/bar/__init__.py", ""),
+        ]);
+        let proj = tmp.path().join("proj");
+        let sp = tmp.path().join("sp");
+        let output = tmp.path().join("out.json");
+
+        let args = RunTreeArgs::try_parse_from([
+            "run-tree",
+            proj.to_str().unwrap(),
+            output.to_str().unwrap(),
+            "--site-packages",
+            sp.to_str().unwrap(),
+            "--sorted-output",
+        ])
+        .unwrap();
+        run(args).unwrap();
+
+        let content = fs::read_to_string(&output).unwrap();
+        let value: Value = serde_json::from_str(&content).unwrap();
+        let modules: BTreeSet<&str> = value["LAZY_ELIGIBLE"]
+            .as_object()
+            .expect("LAZY_ELIGIBLE object in output JSON")
+            .keys()
+            .map(|s| s.as_str())
+            .collect();
+        // `foo` appears only because --site-packages caused its resolution.
+        // `bar` doesn't.
+        assert_eq!(modules, BTreeSet::from(["main", "foo"]));
+    }
+}