Restrict register(...) usage (#151)

* Restrict usage of @register(...) with arguments so that arguments can't be replaced afterwards.

* Move tests specific for `register(...)` to a separate module

* Fix signature

* Restore removed code that breaks mypy

pytest_factoryboy/fixture.py

@@ -1,26 +1,24 @@
 """Factory boy fixture integration."""
 from __future__ import annotations
 
-import functools
 import sys
 from dataclasses import dataclass
 from inspect import signature
-from typing import TYPE_CHECKING, cast, overload
+from typing import TYPE_CHECKING, overload
 
 import factory
 import factory.builder
 import factory.declarations
 import factory.enums
 import inflection
-from typing_extensions import Protocol
 
 from .codegen import FixtureDef, make_fixture_model_module
 from .compat import PostGenerationContext
 
 if TYPE_CHECKING:
     from typing import Any, Callable, TypeVar
 
-    from _pytest.fixtures import FixtureFunction, FixtureRequest, SubRequest
+    from _pytest.fixtures import FixtureFunction, SubRequest
     from factory.builder import BuildStep
     from factory.declarations import PostGeneration, PostGenerationContext
 
@@ -43,41 +41,45 @@ def __call__(self, request: SubRequest) -> Any:
         return self.function(request)
 
 
-class RegisterProtocol(Protocol):
-    """Protocol for ``register`` function called with ``factory_class``."""
-
-    def __call__(self, factory_class: F, _name: str | None = None, **kwargs: Any) -> F:
-        """``register`` function called with ``factory_class``."""
-
-
+# register(AuthorFactory, ...)
+#
+# @register
+# class AuthorFactory(factory.Factory): ...
 @overload
-def register(
-    factory_class: None = None,
-    _name: str | None = None,
-    **kwargs: Any,
-) -> RegisterProtocol:
+def register(factory_class: F, _name: str | None = None, **kwargs: Any) -> F:
     ...
 
 
+# @register(...)
+# class AuthorFactory(factory.Factory): ...
 @overload
-def register(factory_class: F, _name: str | None = None, **kwargs: Any) -> F:
+def register(*, _name: str | None = None, **kwargs: Any) -> Callable[[F], F]:
     ...
 
 
 def register(
     factory_class: F | None = None,
     _name: str | None = None,
+    *,
+    _caller_locals: dict[str, Any] | None = None,
     **kwargs: Any,
-) -> F | RegisterProtocol:
+) -> F | Callable[[F], F]:
     r"""Register fixtures for the factory class.
 
     :param factory_class: Factory class to register.
     :param _name: Name of the model fixture. By default, is lowercase-underscored model name.
+    :param _caller_locals: Dictionary where to inject the generated fixtures. Defaults to the caller's locals().
     :param \**kwargs: Optional keyword arguments that override factory attributes.
     """
+    if _caller_locals is None:
+        _caller_locals = get_caller_locals()
 
     if factory_class is None:
-        return functools.partial(register, _name=_name, **kwargs)
+
+        def register_(factory_class: F) -> F:
+            return register(factory_class, _name=_name, _caller_locals=_caller_locals, **kwargs)
+
+        return register_
 
     assert not factory_class._meta.abstract, "Can't register abstract factories."
     assert factory_class._meta.model is not None, "Factory model class is not specified."
@@ -146,9 +148,7 @@ def register(
                     )
                 )
 
-    caller_locals = get_caller_locals()
-
-    if factory_name not in caller_locals:
+    if factory_name not in _caller_locals:
         fixture_defs.append(
             FixtureDef(
                 name=factory_name,
@@ -172,7 +172,7 @@ def register(
     for fixture_def in fixture_defs:
         exported_name = fixture_def.name
         fixture_function = getattr(generated_module, exported_name)
-        inject_into_caller(exported_name, fixture_function, caller_locals)
+        inject_into_caller(exported_name, fixture_function, _caller_locals)
 
     return factory_class
 

tests/test_factory_fixtures.py

@@ -64,11 +64,6 @@ class Meta:
 
 
 @register
-@register(
-    _name="harry_potter_author",
-    name="J.K. Rowling",
-    register_user="jk_rowling",
-)
 class AuthorFactory(factory.Factory):
     """Author factory."""
 
@@ -161,41 +156,41 @@ def test_post_generation(author: Author):
     assert author.user.is_active is True
 
 
-register(AuthorFactory, "second_author")
-
-
-@pytest.mark.parametrize("second_author__name", ["Mr. Hyde"])
-def test_second_author(author: Author, second_author: Author):
-    """Test factory registration with specific name."""
-    assert author != second_author
-    assert second_author.name == "Mr. Hyde"
+class TestParametrizeAlternativeNameFixture:
+    register(AuthorFactory, "second_author")
 
+    @pytest.mark.parametrize("second_author__name", ["Mr. Hyde"])
+    def test_second_author(self, author: Author, second_author: Author):
+        """Test parametrization of attributes for fixture registered under a different name
+        ("second_author")."""
+        assert author != second_author
+        assert second_author.name == "Mr. Hyde"
 
-register(AuthorFactory, "partial_author", name="John Doe", register_user=LazyFixture(lambda: "jd@jd.com"))
 
+class TestPartialSpecialization:
+    register(AuthorFactory, "partial_author", name="John Doe", register_user=LazyFixture(lambda: "jd@jd.com"))
 
-def test_partial(partial_author: Author):
-    """Test fixture partial specialization."""
-    assert partial_author.name == "John Doe"
-    assert partial_author.user
-    assert partial_author.user.username == "jd@jd.com"
+    def test_partial(self, partial_author: Author):
+        """Test fixture partial specialization."""
+        assert partial_author.name == "John Doe"
+        assert partial_author.user  # Makes mypy happy
+        assert partial_author.user.username == "jd@jd.com"
 
 
-register(AuthorFactory, "another_author", name=LazyFixture(lambda: "Another Author"))
+class TestLazyFixture:
+    register(AuthorFactory, "another_author", name=LazyFixture(lambda: "Another Author"))
 
+    @pytest.mark.parametrize("book__author", [LazyFixture("another_author")])
+    def test_lazy_fixture_name(self, book: Book, another_author: Author):
+        """Test that book author is replaced with another author by fixture name."""
+        assert book.author == another_author
+        assert book.author.name == "Another Author"
 
-@pytest.mark.parametrize("book__author", [LazyFixture("another_author")])
-def test_lazy_fixture_name(book: Book, another_author: Author):
-    """Test that book author is replaced with another author by fixture name."""
-    assert book.author == another_author
-    assert book.author.name == "Another Author"
-
-
-@pytest.mark.parametrize("book__author", [LazyFixture(lambda another_author: another_author)])
-def test_lazy_fixture_callable(book: Book, another_author: Author) -> None:
-    """Test that book author is replaced with another author by callable."""
-    assert book.author == another_author
-    assert book.author.name == "Another Author"
+    @pytest.mark.parametrize("book__author", [LazyFixture(lambda another_author: another_author)])
+    def test_lazy_fixture_callable(self, book: Book, another_author: Author) -> None:
+        """Test that book author is replaced with another author by callable."""
+        assert book.author == another_author
+        assert book.author.name == "Another Author"
 
 
 @pytest.mark.parametrize(
@@ -209,24 +204,3 @@ def test_lazy_fixture_post_generation(author: Author):
     # assert author.user.username == "lazyfixture"
     assert author.user
     assert author.user.password == "asdasd"
-
-
-def test_register_class_decorator_with_kwargs_only(harry_potter_author: Author):
-    """Ensure ``register`` decorator called with kwargs only works normally."""
-    assert harry_potter_author.name == "J.K. Rowling"
-    assert harry_potter_author.user
-    assert harry_potter_author.user.username == "jk_rowling"
-
-
-register(_name="the_chronicles_of_narnia_author", name="C.S. Lewis")(
-    AuthorFactory,
-    register_user="cs_lewis",
-    register_user__password="Aslan1",
-)
-
-
-def test_register_function_with_kwargs_only(the_chronicles_of_narnia_author: Author):
-    """Ensure ``register`` function called with kwargs only works normally."""
-    assert the_chronicles_of_narnia_author.name == "C.S. Lewis"
-    assert the_chronicles_of_narnia_author.user
-    assert the_chronicles_of_narnia_author.user.password == "Aslan1"

tests/test_registration.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import factory
+import pytest
+from _pytest.fixtures import FixtureLookupError
+
+from pytest_factoryboy import register
+
+
+@dataclass(eq=False)
+class Foo:
+    value: str
+
+
+class TestRegisterDirectDecorator:
+    @register()
+    class FooFactory(factory.Factory):
+        class Meta:
+            model = Foo
+
+        value = "@register()"
+
+    def test_register(self, foo: Foo):
+        """Test that `register` can be used as a decorator with 0 arguments."""
+        assert foo.value == "@register()"
+
+
+class TestRegisterDecoratorNoArgs:
+    @register
+    class FooFactory(factory.Factory):
+        class Meta:
+            model = Foo
+
+        value = "@register"
+
+    def test_register(self, foo: Foo):
+        """Test that `register` can be used as a direct decorator."""
+        assert foo.value == "@register"
+
+
+class TestRegisterDecoratorWithArgs:
+    @register(value="bar")
+    class FooFactory(factory.Factory):
+        class Meta:
+            model = Foo
+
+        value = None
+
+    def test_register(self, foo: Foo):
+        """Test that `register` can be used as a decorator with arguments overriding the factory declarations."""
+        assert foo.value == "bar"
+
+
+class TestRegisterAlternativeName:
+    @register(_name="second_foo")
+    class FooFactory(factory.Factory):
+        class Meta:
+            model = Foo
+
+        value = None
+
+    def test_register(self, request, second_foo: Foo):
+        """Test that `register` invoked with a specific `_name` registers the fixture under that `_name`."""
+        assert second_foo.value == None
+
+        with pytest.raises(FixtureLookupError) as exc:
+            request.getfixturevalue("foo")
+        assert exc.value.argname == "foo"
+
+
+class TestRegisterAlternativeNameAndArgs:
+    @register(_name="second_foo", value="second_bar")
+    class FooFactory(factory.Factory):
+        class Meta:
+            model = Foo
+
+        value = None
+
+    def test_register(self, second_foo: Foo):
+        """Test that `register` can be invoked with `_name` to specify an alternative
+        fixture name and with any kwargs to override the factory declarations."""
+        assert second_foo.value == "second_bar"
+
+
+class TestRegisterCall:
+    class FooFactory(factory.Factory):
+        class Meta:
+            model = Foo
+
+        value = "register(FooFactory)"
+
+    register(FooFactory)
+    register(FooFactory, _name="second_foo", value="second_bar")
+
+    def test_register(self, foo: Foo, second_foo: Foo):
+        """Test that `register` can be invoked directly."""
+        assert foo.value == "register(FooFactory)"
+        assert second_foo.value == "second_bar"