fix

Author: asukaminato0721

crates/pyrefly_types/src/function.rs

@@ -242,6 +242,8 @@ pub struct FuncFlags {
     /// A method directly inside a `Protocol` class.
     pub is_in_protocol_class: bool,
     pub is_in_type_checking_block: bool,
+    /// The generated mock type for a callable resolved from a `unittest.mock.patch` target.
+    pub unittest_mock_type: Option<Box<Type>>,
     /// Set when the definition has both `*args` and `**kwargs` and both are typed
     /// `Any` — explicitly, or implicitly because they are unannotated. Per the
     /// typing spec such a signature is equivalent to `...`. Captured at definition

pyrefly/lib/alt/attr.rs

@@ -632,6 +632,9 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         context: Option<&dyn Fn() -> ErrorContext>,
         todo_ctx: &str,
     ) -> Type {
+        if let Some(signature) = self.unittest_mock_assertion_signature(base, attr_name) {
+            return signature;
+        }
         let attr_base = self.as_attribute_base(base.clone());
         let lookup_result = attr_base.clone().map_or_else(
             || LookupResult::internal_error(InternalError::AttributeBaseUndefined(base.clone())),
@@ -721,6 +724,34 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         }
     }
 
+    /// Check call assertions against the signature retained from a `mock.patch` target.
+    fn unittest_mock_assertion_signature(&self, base: &Type, attr_name: &Name) -> Option<Type> {
+        if !matches!(
+            attr_name.as_str(),
+            "assert_called_with"
+                | "assert_called_once_with"
+                | "assert_any_call"
+                | "assert_awaited_with"
+                | "assert_awaited_once_with"
+                | "assert_any_await"
+        ) {
+            return None;
+        }
+        let mock = base
+            .toplevel_func_metadata()?
+            .flags
+            .unittest_mock_type
+            .clone()?;
+        if attr_name.as_str().contains("await")
+            && !matches!(*mock, Type::ClassType(cls) if cls.has_qname("unittest.mock", "AsyncMock"))
+        {
+            return None;
+        }
+        let mut target = base.clone();
+        target.transform_toplevel_callable_signatures(|callable, _| callable.ret = Type::None);
+        Some(target)
+    }
+
     fn add_class_fields(&self, class: &Class, candidates: &mut SmallSet<Name>) {
         let mut add_fields_for = |cls: &Class| {
             if let Some(class_fields) = self.get_class_fields(cls) {
@@ -2587,6 +2618,14 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
     }
 
     fn as_attribute_base1(&self, ty: Type, acc: &mut Vec<AttributeBase1>) {
+        // A patched callable keeps its signature but delegates attributes to the generated mock.
+        if let Some(mock) = ty
+            .toplevel_func_metadata()
+            .and_then(|metadata| metadata.flags.unittest_mock_type.clone())
+        {
+            self.as_attribute_base1(*mock, acc);
+            return;
+        }
         match ty {
             Type::ClassType(class_type) => {
                 if let Some(shaped_array) =

pyrefly/lib/alt/call.rs

@@ -2587,7 +2587,8 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                 )
         } else {
             self.expand_mut(&mut callee_ty);
-            self.check_unittest_mock_patch_target(&callee_ty, &x.arguments, errors);
+            let unittest_mock_patch_target =
+                self.check_unittest_mock_patch_target(&callee_ty, &x.arguments, errors);
 
             let call = CallWithTypes::new();
             let (args, kws) = if callee_ty.is_union() {
@@ -2619,6 +2620,16 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                     );
                 }
                 match ty.callee_kind() {
+                _ if ty.toplevel_func_metadata().is_some_and(|metadata| metadata.flags.unittest_mock_type.is_some()) =>
+                    self.freeform_call_infer(
+                        ty.clone(),
+                        &args,
+                        &kws,
+                        x.func.range(),
+                        x.arguments.range(),
+                        hint,
+                        errors,
+                    ).ty,
                 Some(CalleeKind::Function(FunctionKind::AssertType)) => self
                     .call_assert_type(
                         &x.arguments.args,
@@ -2852,12 +2863,18 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                 }
             }});
             // TypeIs and TypeGuard functions return bool at runtime
-            match result {
+            let result = match result {
                 Type::TypeIs(_) | Type::TypeGuard(_) => {
                     self.heap.mk_class_type(self.stdlib.bool().clone())
                 }
                 other => other,
-            }
+            };
+            let uses_default_mock = x.arguments.find_keyword("new").is_none()
+                && x.arguments.find_keyword("new_callable").is_none();
+            self.specialize_unittest_mock_patch(
+                result,
+                unittest_mock_patch_target.filter(|_| uses_default_mock),
+            )
         };
 
         let result = self.apply_django_annotate_call(result, django_annotate_call);
@@ -2869,33 +2886,31 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         callee_ty: &Type,
         arguments: &Arguments,
         errors: &ErrorCollector,
-    ) {
+    ) -> Option<Type> {
         let Type::ClassType(cls) = callee_ty else {
-            return;
+            return None;
         };
         if !cls.has_qname("unittest.mock", "_patcher") {
-            return;
+            return None;
         }
         if arguments.args.len() > 1 || arguments.keywords.iter().any(|kw| kw.arg.is_none()) {
-            return;
+            return None;
         }
         if arguments
             .find_keyword("create")
             .is_some_and(|kw| !matches!(&kw.value, Expr::BooleanLiteral(lit) if !lit.value))
         {
-            return;
+            return None;
         }
-        let Some(target_expr) = arguments.find_argument_value("target", 0) else {
-            return;
-        };
+        let target_expr = arguments.find_argument_value("target", 0)?;
         let Expr::StringLiteral(ExprStringLiteral { value, .. }) = target_expr else {
-            return;
+            return None;
         };
         let range = target_expr.range();
         let target = value.to_str();
         let parts: Vec<&str> = target.split('.').collect();
         if parts.len() < 2 || parts.iter().any(|p| p.is_empty()) {
-            return;
+            return None;
         }
 
         let Some((module_prefix_len, module)) = (1..parts.len()).rev().find_map(|i| {
@@ -2906,7 +2921,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         }) else {
             // The target may be resolved dynamically at runtime, so only check paths rooted in a
             // module that is available in the current environment.
-            return;
+            return None;
         };
 
         let attrs = &parts[module_prefix_len..];
@@ -2927,7 +2942,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                     .exports
                     .is_implicit_reexport(ModuleName::builtins(), &name)
             {
-                return;
+                return None;
             }
             base_ty = self.type_of_attr_get(
                 &base_ty,
@@ -2939,6 +2954,48 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                 "unittest.mock.patch target",
             );
         }
+        Some(base_ty)
+    }
+
+    /// Retain a resolved patch target's signature on its generated mock result.
+    fn specialize_unittest_mock_patch(&self, result: Type, target: Option<Type>) -> Type {
+        let Some(mut target) = target else {
+            return result;
+        };
+        let Some(metadata) = target.toplevel_func_metadata_mut() else {
+            return result;
+        };
+        let is_async = metadata.flags.is_async;
+        let Type::ClassType(mut patcher) = result else {
+            return result;
+        };
+        if !patcher.has_qname("unittest.mock", "_patch_pass_arg") {
+            return Type::ClassType(patcher);
+        }
+        let [mock] = patcher.targs_mut().as_mut() else {
+            unreachable!("unittest.mock._patch_pass_arg has exactly one type argument")
+        };
+        let Type::Union(mock_union) = mock else {
+            unreachable!("the default unittest.mock.patch result is a union of mock types")
+        };
+        let mock_type = mock_union
+            .members
+            .iter()
+            .find(|ty| {
+                matches!(
+                    ty,
+                    Type::ClassType(cls)
+                        if cls.has_qname(
+                            "unittest.mock",
+                            if is_async { "AsyncMock" } else { "MagicMock" },
+                        )
+                )
+            })
+            .cloned()
+            .expect("the default patch result contains MagicMock and AsyncMock");
+        metadata.flags.unittest_mock_type = Some(Box::new(mock_type));
+        *mock = target;
+        Type::ClassType(patcher)
     }
 
     pub fn freeform_call_infer(

pyrefly/lib/test/unittest_mock_patch.rs

@@ -284,16 +284,15 @@ with patch.multiple("other.missing", attribute=DEFAULT):
 
 // https://github.com/facebook/pyrefly/issues/3913
 testcase!(
-    bug = "Mock assertions do not check the target signature",
     test_assert_called_with_checks_target_signature,
     TestEnv::one("other", "def foo(x: int, *, y: str) -> None: ..."),
     r#"
 from unittest.mock import patch
 
 with patch("other.foo") as mock_foo:
     mock_foo.assert_called_with(1, y="ok")
-    mock_foo.assert_called_once_with("bad", y="ok")  # BUG: This should report that x requires int, not str.
-    mock_foo.assert_any_call(1, y=2)  # BUG: This should report that y requires str, not int.
+    mock_foo.assert_called_once_with("bad", y="ok")  # E: Argument `Literal['bad']` is not assignable to parameter `x` with type `int` in function `other.foo`
+    mock_foo.assert_any_call(1, y=2)  # E: Argument `Literal[2]` is not assignable to parameter `y` with type `str` in function `other.foo`
 "#,
 );
 
@@ -311,36 +310,33 @@ with patch("typing.assert_type") as mock_assert_type:
 
 // https://github.com/facebook/pyrefly/issues/3913
 testcase!(
-    bug = "Mock assertions do not check the target signature",
     test_patch_start_checks_target_signature,
     TestEnv::one("other", "def foo(x: int) -> None: ..."),
     r#"
 from unittest.mock import patch
 
 patcher = patch("other.foo")
 mock_foo = patcher.start()
-mock_foo.assert_called_with("bad")  # BUG: This should report that x requires int, not str.
+mock_foo.assert_called_with("bad")  # E: Argument `Literal['bad']` is not assignable to parameter `x` with type `int` in function `other.foo`
 patcher.stop()
 "#,
 );
 
 // https://github.com/facebook/pyrefly/issues/3913
 testcase!(
-    bug = "Mock assertions do not check the target signature",
     test_async_patch_checks_await_assertion,
     TestEnv::one("other", "async def foo(x: int) -> None: ..."),
     r#"
 from unittest.mock import patch
 
 with patch("other.foo") as mock_foo:
     mock_foo.assert_awaited_with(1)
-    mock_foo.assert_awaited_once_with("bad")  # BUG: This should report that x requires int, not str.
+    mock_foo.assert_awaited_once_with("bad")  # E: Argument `Literal['bad']` is not assignable to parameter `x` with type `int` in function `other.foo`
 "#,
 );
 
 // https://github.com/facebook/pyrefly/issues/3913
 testcase!(
-    bug = "Mock assertions do not check the target signature",
     test_patch_bound_method_checks_bound_signature,
     TestEnv::one(
         "other",
@@ -350,7 +346,7 @@ testcase!(
 from unittest.mock import patch
 
 with patch("other.instance.method") as mock_method:
-    mock_method.assert_called_once_with("bad")  # BUG: This should report that x requires int, not str.
+    mock_method.assert_called_once_with("bad")  # E: Argument `Literal['bad']` is not assignable to parameter `x` with type `int` in function `other.C.method`
 "#,
 );