[Draft] JUnit graph proxy support added

test/test-junit5/src/main/java/ru/tinkoff/kora/test/extension/junit5/GraphProxy.java

@@ -0,0 +1,29 @@
+package ru.tinkoff.kora.test.extension.junit5;
+
+import org.junit.jupiter.api.extension.ExtensionConfigurationException;
+import ru.tinkoff.kora.application.graph.ApplicationGraphDraw;
+import ru.tinkoff.kora.application.graph.Node;
+import ru.tinkoff.kora.application.graph.internal.NodeImpl;
+
+import java.util.function.BiFunction;
+
+record GraphProxy<T>(BiFunction<T, KoraAppGraph, ? extends T> function,
+                     GraphCandidate candidate) implements GraphModification {
+
+    @SuppressWarnings({"rawtypes", "unchecked"})
+    @Override
+    public void accept(ApplicationGraphDraw graphDraw) {
+        var nodesToReplace = GraphUtils.findNodeByTypeOrAssignable(graphDraw, candidate());
+        if (nodesToReplace.isEmpty()) {
+            throw new ExtensionConfigurationException("Can't find Nodes to Proxy: " + candidate());
+        }
+
+        for (var nodeToReplace : nodesToReplace) {
+            var casted = (Node<T>) nodeToReplace;
+            graphDraw.replaceNodeKeepDependencies(casted, g -> {
+                var self = ((NodeImpl) casted).factory.get(g);
+                return function.apply((T) self, new DefaultKoraAppGraph(graphDraw, g));
+            });
+        }
+    }
+}

test/test-junit5/src/main/java/ru/tinkoff/kora/test/extension/junit5/KoraGraphModification.java

@@ -6,6 +6,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
@@ -120,6 +121,94 @@ public <T> KoraGraphModification replaceComponent(@Nonnull Type typeToReplace,
         }
     }
 
+    /**
+     * Component that should replace existing one with new one AND keeps its dependencies in graph, original component is also available
+     */
+    @Nonnull
+    public <T> KoraGraphModification proxyComponent(@Nonnull Type typeToReplace,
+                                                    @Nonnull Function<T, ? extends T> proxyFunction) {
+        modifications.add(new GraphProxy<T>((original, graph) -> proxyFunction.apply(original), new GraphCandidate(typeToReplace)));
+        return this;
+    }
+
+    /**
+     * Component that should replace existing one with new one AND keeps its dependencies in graph, original component is also available
+     */
+    @Nonnull
+    public <T> KoraGraphModification proxyComponent(@Nonnull Type typeToReplace,
+                                                    @Nonnull List<Class<?>> tags,
+                                                    @Nonnull Function<T, ? extends T> proxyFunction) {
+        if (tags.isEmpty()) {
+            return proxyComponent(typeToReplace, proxyFunction);
+        } else {
+            modifications.add(new GraphProxy<T>((original, graph) -> proxyFunction.apply(original), new GraphCandidate(typeToReplace, tags)));
+            return this;
+        }
+    }
+
+    /**
+     * Component that should replace existing one with new one AND keeps its dependencies in graph, original component is also available
+     */
+    @Nonnull
+    public <T> KoraGraphModification proxyComponent(@Nonnull Type typeToReplace,
+                                                    @Nonnull BiFunction<T, KoraAppGraph, ? extends T> proxyFunction) {
+        modifications.add(new GraphProxy<T>(proxyFunction, new GraphCandidate(typeToReplace)));
+        return this;
+    }
+
+    /**
+     * Component that should replace existing one with new one AND keeps its dependencies in graph, original component is also available
+     */
+    @Nonnull
+    public <T> KoraGraphModification proxyComponent(@Nonnull Type typeToReplace,
+                                                    @Nonnull List<Class<?>> tags,
+                                                    @Nonnull BiFunction<T, KoraAppGraph, ? extends T> proxyFunction) {
+        if (tags.isEmpty()) {
+            return proxyComponent(typeToReplace, proxyFunction);
+        } else {
+            modifications.add(new GraphProxy<T>(proxyFunction, new GraphCandidate(typeToReplace, tags)));
+            return this;
+        }
+    }
+
+    /**
+     * Component that should replace existing one with new one AND keeps its dependencies in graph, original component is also available
+     */
+    @Nonnull
+    public <T> KoraGraphModification proxyComponent(@Nonnull Class<T> typeToReplace,
+                                                    @Nonnull Function<T, ? extends T> proxyFunction) {
+        return proxyComponent(((Type) typeToReplace), proxyFunction);
+    }
+
+    /**
+     * Component that should replace existing one with new one AND keeps its dependencies in graph, original component is also available
+     */
+    @Nonnull
+    public <T> KoraGraphModification proxyComponent(@Nonnull Class<T> typeToReplace,
+                                                    @Nonnull List<Class<?>> tags,
+                                                    @Nonnull Function<T, ? extends T> proxyFunction) {
+        return proxyComponent(((Type) typeToReplace), tags, proxyFunction);
+    }
+
+    /**
+     * Component that should replace existing one with new one AND keeps its dependencies in graph, original component is also available
+     */
+    @Nonnull
+    public <T> KoraGraphModification proxyComponent(@Nonnull Class<T> typeToReplace,
+                                                    @Nonnull BiFunction<T, KoraAppGraph, ? extends T> proxyFunction) {
+        return proxyComponent(((Type) typeToReplace), proxyFunction);
+    }
+
+    /**
+     * Component that should replace existing one with new one AND keeps its dependencies in graph, original component is also available
+     */
+    @Nonnull
+    public <T> KoraGraphModification proxyComponent(@Nonnull Class<T> typeToReplace,
+                                                    @Nonnull List<Class<?>> tags,
+                                                    @Nonnull BiFunction<T, KoraAppGraph, ? extends T> proxyFunction) {
+        return proxyComponent(((Type) typeToReplace), tags, proxyFunction);
+    }
+
     /**
      * Component that should replace existing one with Mock AND all real component dependencies removed for graph
      */

test/test-junit5/src/test/java/ru/tinkoff/kora/test/extension/junit5/proxy/ProxyTests.java

@@ -0,0 +1,27 @@
+package ru.tinkoff.kora.test.extension.junit5.proxy;
+
+import jakarta.annotation.Nonnull;
+import org.junit.jupiter.api.Test;
+import ru.tinkoff.kora.test.extension.junit5.KoraAppTest;
+import ru.tinkoff.kora.test.extension.junit5.KoraAppTestGraphModifier;
+import ru.tinkoff.kora.test.extension.junit5.KoraGraphModification;
+import ru.tinkoff.kora.test.extension.junit5.TestComponent;
+import ru.tinkoff.kora.test.extension.junit5.testdata.TestApplication;
+import ru.tinkoff.kora.test.extension.junit5.testdata.TestApplication.SomeFactory;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+@KoraAppTest(TestApplication.class)
+public class ProxyTests implements KoraAppTestGraphModifier {
+
+    @Override
+    public @Nonnull KoraGraphModification graph() {
+        return KoraGraphModification.create()
+            .proxyComponent(SomeFactory.class, (original) -> () -> original.getValue() + "2");
+    }
+
+    @Test
+    void originalWithReplacedBean(@TestComponent SomeFactory someFactory) {
+        assertEquals("12", someFactory.getValue());
+    }
+}

test/test-junit5/src/test/java/ru/tinkoff/kora/test/extension/junit5/proxy/ProxyWithGraphTests.java

@@ -0,0 +1,31 @@
+package ru.tinkoff.kora.test.extension.junit5.proxy;
+
+import jakarta.annotation.Nonnull;
+import org.junit.jupiter.api.Test;
+import ru.tinkoff.kora.test.extension.junit5.KoraAppTest;
+import ru.tinkoff.kora.test.extension.junit5.KoraAppTestGraphModifier;
+import ru.tinkoff.kora.test.extension.junit5.KoraGraphModification;
+import ru.tinkoff.kora.test.extension.junit5.TestComponent;
+import ru.tinkoff.kora.test.extension.junit5.testdata.TestApplication;
+import ru.tinkoff.kora.test.extension.junit5.testdata.TestApplication.SomeFactory;
+import ru.tinkoff.kora.test.extension.junit5.testdata.TestComponent1;
+import ru.tinkoff.kora.test.extension.junit5.testdata.TestComponent12;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+@KoraAppTest(value = TestApplication.class, components = TestComponent12.class)
+public class ProxyWithGraphTests implements KoraAppTestGraphModifier {
+
+    @Override
+    public @Nonnull KoraGraphModification graph() {
+        return KoraGraphModification.create()
+            .proxyComponent(SomeFactory.class, (original, graph) -> {
+                return () -> original.getValue() + "2" + graph.getFirst(TestComponent12.class).get();
+            });
+    }
+
+    @Test
+    void originalWithReplacedBean(@TestComponent SomeFactory someFactory) {
+        assertEquals("1212", someFactory.getValue());
+    }
+}

test/test-junit5/src/test/java/ru/tinkoff/kora/test/extension/junit5/proxy/ProxyWithTagTests.java

@@ -0,0 +1,32 @@
+package ru.tinkoff.kora.test.extension.junit5.proxy;
+
+import jakarta.annotation.Nonnull;
+import org.junit.jupiter.api.Test;
+import ru.tinkoff.kora.common.Tag;
+import ru.tinkoff.kora.test.extension.junit5.KoraAppTest;
+import ru.tinkoff.kora.test.extension.junit5.KoraAppTestGraphModifier;
+import ru.tinkoff.kora.test.extension.junit5.KoraGraphModification;
+import ru.tinkoff.kora.test.extension.junit5.TestComponent;
+import ru.tinkoff.kora.test.extension.junit5.testdata.TestApplication;
+import ru.tinkoff.kora.test.extension.junit5.testdata.TestApplication.SomeFactory;
+
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+@KoraAppTest(TestApplication.class)
+public class ProxyWithTagTests implements KoraAppTestGraphModifier {
+
+    @Override
+    public @Nonnull KoraGraphModification graph() {
+        return KoraGraphModification.create()
+            .proxyComponent(SomeFactory.class, List.of(String.class), (original) -> {
+                return () -> original.getValue() + "2";
+            });
+    }
+
+    @Test
+    void originalWithReplacedBean(@Tag(String.class) @TestComponent SomeFactory someFactory) {
+        assertEquals("12", someFactory.getValue());
+    }
+}

test/test-junit5/src/test/java/ru/tinkoff/kora/test/extension/junit5/testdata/TestApplication.java

@@ -3,6 +3,7 @@
 import ru.tinkoff.kora.application.graph.LifecycleWrapper;
 import ru.tinkoff.kora.application.graph.Wrapped;
 import ru.tinkoff.kora.common.KoraApp;
+import ru.tinkoff.kora.common.Tag;
 import ru.tinkoff.kora.common.annotation.Root;
 
 import java.util.function.Function;
@@ -49,6 +50,17 @@ default Function<String, Integer> consumerExample() {
         return (s) -> 1;
     }
 
+    @Root
+    default SomeFactory someFactory() {
+        return () -> "1";
+    }
+
+    @Root
+    @Tag(String.class)
+    default SomeFactory someFactoryWithTag() {
+        return () -> "1";
+    }
+
     class CustomWrapper implements Wrapped<Float> {
 
         @Override
@@ -57,6 +69,11 @@ public Float value() {
         }
     }
 
+    interface SomeFactory {
+
+        String getValue();
+    }
+
     interface SomeParent {
 
     }