From 3785581c6542112ea4d767ecf86b89a453400241 Mon Sep 17 00:00:00 2001
From: Mats <mats@craftmend.com>
Date: Thu, 24 Aug 2023 17:35:05 +0200
Subject: [PATCH] Fix test case

---
 rpcnis-core/src/main/java/com/rpcnis/core/Rpcnis.java    | 9 +++++++--
 .../java/com/rpcnis/core/proxy/ProxyInvocHandler.java    | 2 +-
 .../npcnis/cor/invocations/PendingInvocationTest.java    | 3 +--
 3 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/rpcnis-core/src/main/java/com/rpcnis/core/Rpcnis.java b/rpcnis-core/src/main/java/com/rpcnis/core/Rpcnis.java
index 3d6898a..ecb7953 100644
--- a/rpcnis-core/src/main/java/com/rpcnis/core/Rpcnis.java
+++ b/rpcnis-core/src/main/java/com/rpcnis/core/Rpcnis.java
@@ -10,6 +10,7 @@
 
 import java.util.Timer;
 import java.util.UUID;
+import java.util.concurrent.CompletionException;
 import java.util.concurrent.ConcurrentHashMap;
 
 public class Rpcnis {
@@ -51,7 +52,7 @@ public Rpcnis(RpcTransport transport) {
         this(new RpcOptions(), new GsonSerializer(), transport);
     }
 
-    public <T> T invoke(InvocationDescriptor invocationDescriptor, Class<T> returnType) throws InvocationTimedOutException {
+    public <T> T invoke(InvocationDescriptor invocationDescriptor, Class<T> returnType) throws Throwable {
         // create a pending invocation
         PendingInvocation<T> pendingInvocation = new PendingInvocation<>(this, invocationDescriptor, () -> {
             // remove the pending invocation from the map
@@ -64,7 +65,11 @@ public <T> T invoke(InvocationDescriptor invocationDescriptor, Class<T> returnTy
         // TODO: transmit
 
         // wait for response or timeout
-        return pendingInvocation.waitForResponse();
+        try {
+            return pendingInvocation.waitForResponse();
+        } catch (CompletionException e) {
+            throw e.getCause();
+        }
     }
 
     public void completeInvocation(InvocationDescriptor invocationDescriptor, Object value) {
diff --git a/rpcnis-core/src/main/java/com/rpcnis/core/proxy/ProxyInvocHandler.java b/rpcnis-core/src/main/java/com/rpcnis/core/proxy/ProxyInvocHandler.java
index 61f1659..35c0b33 100644
--- a/rpcnis-core/src/main/java/com/rpcnis/core/proxy/ProxyInvocHandler.java
+++ b/rpcnis-core/src/main/java/com/rpcnis/core/proxy/ProxyInvocHandler.java
@@ -18,7 +18,7 @@ public ProxyInvocHandler(Rpcnis rpcnis, String targetName) {
     }
 
     @Override
-    public Object invoke(Object proxy, Method method, Object[] args) throws InvocationTimedOutException {
+    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
         // build invocation descriptor
         Class<?>[] argTypes = new Class<?>[args.length];
         for (int i = 0; i < args.length; i++)
diff --git a/rpcnis-core/src/test/java/com/npcnis/cor/invocations/PendingInvocationTest.java b/rpcnis-core/src/test/java/com/npcnis/cor/invocations/PendingInvocationTest.java
index d0945df..c300df0 100644
--- a/rpcnis-core/src/test/java/com/npcnis/cor/invocations/PendingInvocationTest.java
+++ b/rpcnis-core/src/test/java/com/npcnis/cor/invocations/PendingInvocationTest.java
@@ -28,7 +28,7 @@ public static void tearDown() {
     }
 
     @Test
-    public void testPendingInvocation() {
+    public void testPendingInvocation() throws Throwable {
         // base instance
         Rpcnis rpcnis = new Rpcnis(new LoopbackTransport());
 
@@ -47,7 +47,6 @@ public void testPendingInvocation() {
         });
 
         String response = rpcnis.invoke(invocationDescriptor, String.class);
-
         assert response.equals(testString);
     }