From cf95e74076eddd4a32109bdc9ee71f64e7266b4b Mon Sep 17 00:00:00 2001
From: Lev Nadeinsky <87647282+lev-n@users.noreply.github.com>
Date: Fri, 6 Jun 2025 12:07:49 +0200
Subject: [PATCH] Add unit tests for LogInjectionAgent

---
 pom.xml                                  | 14 ++++++
 src/test/java/LogInjectionAgentTest.java | 54 ++++++++++++++++++++++++
 2 files changed, 68 insertions(+)
 create mode 100644 src/test/java/LogInjectionAgentTest.java

diff --git a/pom.xml b/pom.xml
index 1614841..651a494 100644
--- a/pom.xml
+++ b/pom.xml
@@ -27,6 +27,12 @@
             <artifactId>byte-buddy-agent</artifactId>
             <version>1.14.10</version>
         </dependency>
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter</artifactId>
+            <version>5.10.2</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>
@@ -56,6 +62,14 @@
                     </execution>
                 </executions>
             </plugin>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-surefire-plugin</artifactId>
+                <version>3.0.0</version>
+                <configuration>
+                    <useModulePath>false</useModulePath>
+                </configuration>
+            </plugin>
         </plugins>
     </build>
 </project>
diff --git a/src/test/java/LogInjectionAgentTest.java b/src/test/java/LogInjectionAgentTest.java
new file mode 100644
index 0000000..ef710b4
--- /dev/null
+++ b/src/test/java/LogInjectionAgentTest.java
@@ -0,0 +1,54 @@
+import net.bytebuddy.agent.ByteBuddyAgent;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.ByteArrayOutputStream;
+import java.io.PrintStream;
+import java.lang.instrument.Instrumentation;
+import java.lang.reflect.Field;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+public class LogInjectionAgentTest {
+
+    private Instrumentation getInstrumentation() {
+        return ByteBuddyAgent.install();
+    }
+
+    @AfterEach
+    public void tearDown() throws Exception {
+        // Stop the HTTP server and reset static fields
+        Field serverField = LogInjectionAgent.class.getDeclaredField("apiServer");
+        serverField.setAccessible(true);
+        Object server = serverField.get(null);
+        if (server != null) {
+            server.getClass().getMethod("stop").invoke(server);
+        }
+        Field managerField = LogInjectionAgent.class.getDeclaredField("logPointManager");
+        managerField.setAccessible(true);
+        managerField.set(null, null);
+        serverField.set(null, null);
+    }
+
+    @Test
+    public void testPremainInitializesManager() {
+        Instrumentation inst = getInstrumentation();
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+        PrintStream original = System.out;
+        System.setOut(new PrintStream(out));
+        try {
+            LogInjectionAgent.premain("apiPort=0", inst);
+        } finally {
+            System.setOut(original);
+        }
+        assertTrue(out.toString().contains("Starting Log Injection Agent"));
+        assertNotNull(LogInjectionAgent.getLogPointManager());
+    }
+
+    @Test
+    public void testAgentmainDelegatesToPremain() {
+        Instrumentation inst = getInstrumentation();
+        LogInjectionAgent.agentmain("apiPort=0", inst);
+        assertNotNull(LogInjectionAgent.getLogPointManager());
+    }
+}