Add tests for calling a custom annotate function that only implements VALUE annotations

Author: DavidCEllis

Lib/test/test_annotationlib.py

@@ -10,6 +10,7 @@
 from string.templatelib import Template
 import typing
 import unittest
+import unittest.mock
 from annotationlib import (
     Format,
     ForwardRef,
@@ -1206,6 +1207,74 @@ def evaluate(format, exc=NotImplementedError):
         )
 
 
+class TestCallAnnotateFunction(unittest.TestCase):
+    def _annotate_mock(self):
+        def annotate(format, /):
+            if format == Format.VALUE:
+                return {"x": str}
+            else:
+                raise NotImplementedError(format)
+
+        annotate_mock = unittest.mock.MagicMock(
+            wraps=annotate
+        )
+
+        # Add missing magic attributes needed
+        required_magic = [
+            "__builtins__",
+            "__closure__",
+            "__code__",
+            "__defaults__",
+            "__globals__",
+            "__kwdefaults__",
+        ]
+
+        for attrib in required_magic:
+            setattr(annotate_mock, attrib, getattr(annotate, attrib))
+
+        return annotate_mock
+
+    def test_user_annotate_value(self):
+        annotate = self._annotate_mock()
+
+        annotations = annotationlib.call_annotate_function(
+            annotate,
+            Format.VALUE,
+        )
+
+        self.assertEqual(annotations, {"x": str})
+        annotate.assert_called_once_with(Format.VALUE)
+
+    def test_user_annotate_forwardref(self):
+        annotate = self._annotate_mock()
+
+        with self.assertRaises(NotImplementedError):
+            annotations = annotationlib.call_annotate_function(
+                annotate,
+                Format.FORWARDREF,
+            )
+
+        annotate.assert_has_calls([
+            unittest.mock.Call(Format.FORWARDREF),
+            unittest.mock.Call(Format.VALUE_WITH_FAKE_GLOBALS),
+        ])
+
+    def test_user_annotate_string(self):
+        annotate = self._annotate_mock()
+
+        with self.assertRaises(NotImplementedError):
+            annotations = annotationlib.call_annotate_function(
+                annotate,
+                Format.STRING,
+            )
+
+        annotate.assert_has_calls([
+            unittest.mock.Call(Format.STRING),
+            unittest.mock.Call(Format.VALUE_WITH_FAKE_GLOBALS),
+        ])
+
+
+
 class MetaclassTests(unittest.TestCase):
     def test_annotated_meta(self):
         class Meta(type):