postfetch test

validity/tests/test_utils/test_orm.py

@@ -4,8 +4,9 @@
 from django.db.models import BigIntegerField
 from django.db.models.fields.json import KeyTextTransform
 from django.db.models.functions import Cast
-from factories import DeviceFactory, SerializerDBFactory
+from factories import DeviceFactory, NameSetDBFactory, SerializerDBFactory
 
+from validity.models import NameSet
 from validity.models.serializer import Serializer
 from validity.utils.orm import CustomPrefetchMixin, QuerySetMap
 
@@ -38,3 +39,20 @@ def test_custom_prefetch():
         serializer_id=Cast(KeyTextTransform("serializer", "custom_field_data"), BigIntegerField())
     ).custom_prefetch("serializer", Serializer.objects.all()):
         assert device_serializer_map[device.pk] == device.serializer.pk
+
+
+@pytest.mark.django_db
+def test_custom_postfetch(monkeypatch):
+    DeviceFactory(name="dev0")
+    DeviceFactory(name="dev1")
+    DeviceFactory(name="dev2")
+    device_qs = Device.objects.all()
+    Device.ns = property(lambda self: "ns" + self.name[-1])
+    custom_qs = CustomPrefetchMixin(
+        device_qs.model, device_qs._query, device_qs._db, device_qs._hints
+    ).custom_postfetch("nameset", NameSet.objects.all(), pk_field="ns", remote_pk_field="name")
+    for i in range(4):
+        NameSetDBFactory(name=f"ns{i}")
+    for device in custom_qs:
+        assert device.name.replace("dev", "ns") == device.nameset.name
+    del Device.ns

validity/utils/orm.py

@@ -84,7 +84,7 @@ def get_prefetch_qs_map(self, main_queryset: QuerySet) -> QuerySetMap:
         return self._qs_map_from_keys(pk_values)
 
     def get_postfetch_qs_map(self, main_queryset: QuerySet):
-        pk_values = (getattr(obj, self.remote_pk_field) for obj in main_queryset._result_cache)
+        pk_values = (getattr(obj, self.pk_field) for obj in main_queryset._result_cache)
         return self._qs_map_from_keys(pk_values)
 
     def get_qs_map(self, main_queryset: QuerySet):