[flink] 04/05: [hotfix] Add ExceptionUtils#stripException

[tr...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

trohrmann pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit 5d291961604be26f8be1228e2bf68ab39eec30ed
Author: Till Rohrmann <tr...@apache.org>
AuthorDate: Wed Sep 19 17:28:19 2018 +0200

    [hotfix] Add ExceptionUtils#stripException
    
    stripException strips a given throwable from a specified exception type. This
    is useful to unwrap exceptions.
---
 .../java/org/apache/flink/util/ExceptionUtils.java | 24 ++++++++++++++--------
 .../org/apache/flink/util/ExceptionUtilsTest.java  | 18 ++++++++++++++++
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/flink-core/src/main/java/org/apache/flink/util/ExceptionUtils.java b/flink-core/src/main/java/org/apache/flink/util/ExceptionUtils.java
index 601a252..0ea43ee 100644
--- a/flink-core/src/main/java/org/apache/flink/util/ExceptionUtils.java
+++ b/flink-core/src/main/java/org/apache/flink/util/ExceptionUtils.java
@@ -382,11 +382,7 @@ public final class ExceptionUtils {
 	 * @return Cause of ExecutionException or given Throwable
 	 */
 	public static Throwable stripExecutionException(Throwable throwable) {
-		while (throwable instanceof ExecutionException && throwable.getCause() != null) {
-			throwable = throwable.getCause();
-		}
-
-		return throwable;
+		return stripException(throwable, ExecutionException.class);
 	}
 
 	/**
@@ -397,11 +393,23 @@ public final class ExceptionUtils {
 	 * @return Cause of CompletionException or given Throwable
 	 */
 	public static Throwable stripCompletionException(Throwable throwable) {
-		while (throwable instanceof CompletionException && throwable.getCause() != null) {
-			throwable = throwable.getCause();
+		return stripException(throwable, CompletionException.class);
+	}
+
+	/**
+	 * Unpacks an specified exception and returns its cause. Otherwise the given
+	 * {@link Throwable} is returned.
+	 *
+	 * @param throwableToStrip to strip
+	 * @param typeToStrip type to strip
+	 * @return Unpacked cause or given Throwable if not packed
+	 */
+	public static Throwable stripException(Throwable throwableToStrip, Class<? extends Throwable> typeToStrip) {
+		while (typeToStrip.isAssignableFrom(throwableToStrip.getClass()) && throwableToStrip.getCause() != null) {
+			throwableToStrip = throwableToStrip.getCause();
 		}
 
-		return throwable;
+		return throwableToStrip;
 	}
 
 	/**
diff --git a/flink-core/src/test/java/org/apache/flink/util/ExceptionUtilsTest.java b/flink-core/src/test/java/org/apache/flink/util/ExceptionUtilsTest.java
index 07978a5..fa275d0 100644
--- a/flink-core/src/test/java/org/apache/flink/util/ExceptionUtilsTest.java
+++ b/flink-core/src/test/java/org/apache/flink/util/ExceptionUtilsTest.java
@@ -20,9 +20,12 @@ package org.apache.flink.util;
 
 import org.junit.Test;
 
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.is;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -69,4 +72,19 @@ public class ExceptionUtilsTest extends TestLogger {
 			IllegalStateException.class).isPresent());
 	}
 
+	@Test
+	public void testExceptionStripping() {
+		final FlinkException expectedException = new FlinkException("test exception");
+		final Throwable strippedException = ExceptionUtils.stripException(new RuntimeException(new RuntimeException(expectedException)), RuntimeException.class);
+
+		assertThat(strippedException, is(equalTo(expectedException)));
+	}
+
+	@Test
+	public void testInvalidExceptionStripping() {
+		final FlinkException expectedException = new FlinkException(new RuntimeException(new FlinkException("inner exception")));
+		final Throwable strippedException = ExceptionUtils.stripException(expectedException, RuntimeException.class);
+
+		assertThat(strippedException, is(equalTo(expectedException)));
+	}
 }