Add configuration methods for entity agent factory

junit-common/src/main/java/org/jboss/weld/junit/AbstractWeldInitiator.java

@@ -91,13 +91,16 @@ public static Weld createDefaultWeld() {
 
     private final Function<InjectionPoint, Object> persistenceContextFactory;
 
+    private final Function<InjectionPoint, Object> persistenceAgentFactory;
+
     protected volatile WeldContainer container;
 
     protected AbstractWeldInitiator(Weld weld, List<Object> instancesToInject,
             Set<Class<? extends Annotation>> scopesToActivate, Set<Bean<?>> beans,
             Map<String, Object> resources, Function<InjectionPoint, Object> ejbFactory,
             Function<InjectionPoint, Object> persistenceUnitFactory,
-            Function<InjectionPoint, Object> persistenceContextFactory) {
+            Function<InjectionPoint, Object> persistenceContextFactory,
+            Function<InjectionPoint, Object> persistenceAgentFactory) {
         this.instancesToInject = new ArrayList<>();
         for (Object instance : instancesToInject) {
             this.instancesToInject.add(createToInject(instance));
@@ -137,6 +140,7 @@ protected AbstractWeldInitiator(Weld weld, List<Object> instancesToInject,
         this.ejbFactory = ejbFactory;
         this.persistenceContextFactory = persistenceContextFactory;
         this.persistenceUnitFactory = persistenceUnitFactory;
+        this.persistenceAgentFactory = persistenceAgentFactory;
     }
 
     protected ToInject createToInject(Object instanceToInject) {
@@ -359,6 +363,8 @@ protected static abstract class AbstractBuilder<I extends AbstractWeldInitiator,
 
         private Function<InjectionPoint, Object> persistenceContextFactory;
 
+        private Function<InjectionPoint, Object> persistenceAgentFactory;
+
         public AbstractBuilder(Weld weld) {
             this.weld = weld;
             this.instancesToInject = new ArrayList<>();
@@ -404,6 +410,10 @@ protected Function<InjectionPoint, Object> getPersistenceUnitFactory() {
             return persistenceUnitFactory;
         }
 
+        protected Function<InjectionPoint, Object> getPersistenceAgentFactory() {
+            return persistenceAgentFactory;
+        }
+
         /**
          * Instructs the initiator to inject the given non-contextual instance once the container is started, i.e. during test
          * execution.
@@ -519,6 +529,18 @@ public T setPersistenceContextFactory(Function<InjectionPoint, Object> persisten
             return self();
         }
 
+        /**
+         * Makes it possible to mock {@code PersistenceAgent} injection points.
+         *
+         * @param persistenceAgentFactory
+         * @return self
+         * @since 6.0
+         */
+        public T setPersistenceAgentFactory(Function<InjectionPoint, Object> persistenceAgentFactory) {
+            this.persistenceAgentFactory = persistenceAgentFactory;
+            return self();
+        }
+
         protected abstract T self();
 
         protected abstract I build(Weld weld, List<Object> instancesToInject, Set<Class<? extends Annotation>> scopesToActivate,
@@ -545,8 +567,9 @@ protected WeldContainer initWeldContainer(Weld weld) {
         if (ejbFactory != null) {
             weld.addServices(new MockEjbInjectionServices(ejbFactory));
         }
-        if (persistenceContextFactory != null || persistenceUnitFactory != null) {
-            weld.addServices(new MockJpaInjectionServices(persistenceUnitFactory, persistenceContextFactory));
+        if (persistenceContextFactory != null || persistenceUnitFactory != null || persistenceAgentFactory != null) {
+            weld.addServices(
+                    new MockJpaInjectionServices(persistenceUnitFactory, persistenceContextFactory, persistenceAgentFactory));
         }
         // Init the container
         container = weld.initialize();

junit-common/src/main/java/org/jboss/weld/junit/MockJpaInjectionServices.java

@@ -19,6 +19,7 @@
 import java.util.function.Function;
 
 import jakarta.enterprise.inject.spi.InjectionPoint;
+import jakarta.persistence.EntityAgent;
 import jakarta.persistence.EntityManager;
 import jakarta.persistence.EntityManagerFactory;
 
@@ -39,10 +40,14 @@ public class MockJpaInjectionServices implements JpaInjectionServices {
 
     private final Function<InjectionPoint, Object> persistenceContextFactory;
 
+    private final Function<InjectionPoint, Object> persistenceAgentFactory;
+
     public MockJpaInjectionServices(Function<InjectionPoint, Object> persistenceUnitFactory,
-            Function<InjectionPoint, Object> persistenceContextFactory) {
+            Function<InjectionPoint, Object> persistenceContextFactory,
+            Function<InjectionPoint, Object> persistenceAgentFactory) {
         this.persistenceUnitFactory = persistenceUnitFactory;
         this.persistenceContextFactory = persistenceContextFactory;
+        this.persistenceAgentFactory = persistenceAgentFactory;
     }
 
     @Override
@@ -81,6 +86,24 @@ public ResourceReference<EntityManagerFactory> createResource() {
         };
     }
 
+    @Override
+    public ResourceReferenceFactory<EntityAgent> registerPersistenceAgentInjectionPoint(InjectionPoint injectionPoint) {
+        return new ResourceReferenceFactory<EntityAgent>() {
+            @Override
+            public ResourceReference<EntityAgent> createResource() {
+                if (persistenceAgentFactory == null) {
+                    throw new IllegalStateException(
+                            "Persistence agent factory not set, cannot resolve injection point: " + injectionPoint);
+                }
+                Object agent = persistenceAgentFactory.apply(injectionPoint);
+                if (agent == null || agent instanceof EntityAgent) {
+                    return new SimpleResourceReference<EntityAgent>((EntityAgent) agent);
+                }
+                throw new IllegalStateException("Not an EntityAgent instance: " + agent);
+            }
+        };
+    }
+
     @Override
     public void cleanup() {
     }

junit-jupiter/src/main/java/org/jboss/weld/junit/jupiter/WeldInitiator.java

@@ -182,7 +182,7 @@ protected Builder self() {
         protected WeldInitiator build(Weld weld, List<Object> instancesToInject,
                 Set<Class<? extends Annotation>> scopesToActivate, Set<Bean<?>> beans) {
             return new WeldInitiator(weld, instancesToInject, scopesToActivate, beans, resources, getEjbFactory(),
-                    getPersistenceUnitFactory(), getPersistenceContextFactory());
+                    getPersistenceUnitFactory(), getPersistenceContextFactory(), getPersistenceAgentFactory());
         }
 
     }
@@ -191,9 +191,10 @@ private WeldInitiator(Weld weld, List<Object> instancesToInject, Set<Class<? ext
             Set<Bean<?>> beans,
             Map<String, Object> resources, Function<InjectionPoint, Object> ejbFactory,
             Function<InjectionPoint, Object> persistenceUnitFactory,
-            Function<InjectionPoint, Object> persistenceContextFactory) {
+            Function<InjectionPoint, Object> persistenceContextFactory,
+            Function<InjectionPoint, Object> persistenceAgentFactory) {
         super(weld, instancesToInject, scopesToActivate, beans, resources, ejbFactory, persistenceUnitFactory,
-                persistenceContextFactory);
+                persistenceContextFactory, persistenceAgentFactory);
     }
 
     void shutdownWeld() {

junit-jupiter/src/test/java/org/jboss/weld/junit/jupiter/resources/FooJpa.java

@@ -17,8 +17,10 @@
 package org.jboss.weld.junit.jupiter.resources;
 
 import jakarta.enterprise.context.Dependent;
+import jakarta.persistence.EntityAgent;
 import jakarta.persistence.EntityManager;
 import jakarta.persistence.EntityManagerFactory;
+import jakarta.persistence.PersistenceAgent;
 import jakarta.persistence.PersistenceContext;
 import jakarta.persistence.PersistenceUnit;
 
@@ -31,4 +33,7 @@ public class FooJpa {
     @PersistenceUnit
     EntityManagerFactory entityManagerFactory;
 
+    @PersistenceAgent
+    EntityAgent entityAgent;
+
 }