Add tests for coverage check and coverage --prefer-stubs

pyrefly/lib/commands/coverage/check.rs

@@ -117,3 +117,65 @@ impl CheckArgs {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use pyrefly_util::thread_pool::TEST_THREAD_COUNT;
+    use tempfile::TempDir;
+
+    use super::*;
+
+    /// Run `pyrefly coverage check` on a one-file project with the given extra args.
+    fn run_check(source: &str, extra_args: &[&str]) -> anyhow::Result<CommandExitStatus> {
+        let dir = TempDir::new().unwrap();
+        std::fs::write(
+            dir.path().join("pyrefly.toml"),
+            "search-path = ['.']\nskip-interpreter-query = true\n",
+        )
+        .unwrap();
+        std::fs::write(dir.path().join("main.py"), source).unwrap();
+        let file = dir.path().join("main.py").display().to_string();
+        let args = ["coverage-check", file.as_str()]
+            .into_iter()
+            .chain(extra_args.iter().copied());
+        CheckArgs::parse_from(args).run(None, TEST_THREAD_COUNT)
+    }
+
+    /// Coverage exactly at the threshold passes; just below it exits with a user error.
+    #[test]
+    fn test_fail_under_boundary() {
+        let source = "def f(x) -> int: ...\n"; // 1 of 2 slots typed
+        let check = |threshold| run_check(source, &["--fail-under", threshold]).unwrap();
+        assert_eq!(check("50"), CommandExitStatus::Success);
+        assert_eq!(check("50.01"), CommandExitStatus::UserError);
+    }
+
+    /// A fully-typed project passes the default 100% threshold.
+    #[test]
+    fn test_fully_typed_passes() {
+        assert_eq!(
+            run_check("def f(x: int) -> int: ...\n", &[]).unwrap(),
+            CommandExitStatus::Success
+        );
+    }
+
+    /// `Any` annotations count as covered by default, but not under `--strict` (gh-4024).
+    #[test]
+    fn test_strict_any() {
+        let source = "from typing import Any\n\ndef f(x: Any) -> Any: ...\n";
+        assert_eq!(run_check(source, &[]).unwrap(), CommandExitStatus::Success);
+        assert_eq!(
+            run_check(source, &["--strict"]).unwrap(),
+            CommandExitStatus::UserError
+        );
+    }
+
+    /// Out-of-range `--fail-under` values are rejected before any checking happens.
+    #[test]
+    fn test_fail_under_out_of_range() {
+        for arg in ["--fail-under=100.1", "--fail-under=-0.1"] {
+            let err = run_check("", &[arg]).unwrap_err();
+            assert!(err.to_string().contains("--fail-under"), "{err}");
+        }
+    }
+}

pyrefly/lib/commands/coverage/collect.rs

@@ -3182,4 +3182,98 @@ def g(x: int) -> int:
         let report = build_module_report_for_test_with_env("schema_classes_attrs.py", env);
         compare_snapshot("schema_classes_attrs.expected.json", &report);
     }
+
+    // --- prefer-stubs end-to-end tests ---
+
+    const STUB_TEST_CONFIG: &str = "search-path = ['.']\nskip-interpreter-query = true\n";
+    /// `g` exists only in the `.py`, untyped.
+    const STUB_TEST_PY: &str = "def f() -> int: ...\ndef g(x): ...\n";
+    const STUB_TEST_PYI: &str = "def f() -> int: ...\n";
+
+    /// Run `collect_module_reports` over `dir` with `prefer_stubs` enabled.
+    fn collect_reports_with_stubs(dir: &Path) -> (Vec<ModuleReport>, Vec<Error>) {
+        let globs = FilteredGlobs::new(
+            Globs::new(vec![dir.display().to_string()]).unwrap(),
+            Globs::empty(),
+            None,
+            HiddenDirFilter::Disabled,
+        );
+        collect_module_reports(
+            Box::new(globs),
+            default_config_finder(None),
+            true,
+            None,
+            false,
+            Some(false),
+            TEST_THREAD_COUNT,
+        )
+        .unwrap()
+    }
+
+    /// One stub report: `f` typed, `g` merged as untyped with its finding in the `.py`.
+    fn assert_stub_merged_report(reports: &[ModuleReport], errors: &[Error]) {
+        let [report] = reports else {
+            panic!("expected exactly one report, got {}", reports.len());
+        };
+        assert_eq!(report.name, "foo");
+        assert!(report.path.ends_with("foo.pyi"), "{}", report.path);
+
+        let slots = |name: &str| {
+            report
+                .symbol_reports
+                .iter()
+                .find(|s| s.name() == name)
+                .unwrap_or_else(|| panic!("no symbol named {name}"))
+                .slots()
+        };
+        assert_eq!((slots("foo.f").n_typable, slots("foo.f").n_typed), (1, 1));
+        assert_eq!((slots("foo.g").n_typable, slots("foo.g").n_untyped), (2, 2));
+
+        let [error] = errors else {
+            panic!("expected exactly one finding, got {}", errors.len());
+        };
+        assert!(
+            error.msg_header().contains("foo.g"),
+            "{}",
+            error.msg_header()
+        );
+        assert!(
+            error.path().to_string().ends_with("foo.py"),
+            "{}",
+            error.path()
+        );
+    }
+
+    /// A co-located `.pyi` shadows its `.py`; uncovered `.py` symbols merge into its report.
+    #[test]
+    fn test_prefer_stubs_end_to_end() {
+        let dir = TempDir::new().unwrap();
+        std::fs::write(dir.path().join("pyrefly.toml"), STUB_TEST_CONFIG).unwrap();
+        std::fs::write(dir.path().join("foo.py"), STUB_TEST_PY).unwrap();
+        std::fs::write(dir.path().join("foo.pyi"), STUB_TEST_PYI).unwrap();
+
+        let (reports, errors) = collect_reports_with_stubs(dir.path());
+        assert_stub_merged_report(&reports, &errors);
+    }
+
+    /// A stubs-only project's `.py` is found via site-package-path and merged the same way.
+    #[test]
+    fn test_prefer_stubs_site_package_fallback() {
+        let site = TempDir::new().unwrap();
+        std::fs::write(site.path().join("foo.py"), STUB_TEST_PY).unwrap();
+
+        let dir = TempDir::new().unwrap();
+        std::fs::write(
+            dir.path().join("pyrefly.toml"),
+            format!(
+                "{STUB_TEST_CONFIG}site-package-path = ['{}']\n",
+                site.path().display()
+            ),
+        )
+        .unwrap();
+        std::fs::write(dir.path().join("foo.pyi"), STUB_TEST_PYI).unwrap();
+
+        let (reports, errors) = collect_reports_with_stubs(dir.path());
+        assert_stub_merged_report(&reports, &errors);
+    }
 }

pyrefly/lib/commands/coverage/report.rs

@@ -87,6 +87,8 @@ impl ReportArgs {
 
 #[cfg(test)]
 mod tests {
+    use pyrefly_util::thread_pool::TEST_THREAD_COUNT;
+
     use super::*;
 
     #[test]
@@ -103,4 +105,12 @@ mod tests {
         assert!(parts[0].parse::<u32>().is_ok());
         assert!(parts[1].parse::<u32>().is_ok());
     }
+
+    /// `--module` and `--public-only` cannot be combined.
+    #[test]
+    fn test_module_with_public_only_rejected() {
+        let args = ReportArgs::parse_from(["report", "--module", "m", "--public-only"]);
+        let err = args.run(None, TEST_THREAD_COUNT).unwrap_err();
+        assert!(err.to_string().contains("cannot be combined"), "{err}");
+    }
 }