From 89f59113927dd886f09d8fe2c05ff2cd5d1390c1 Mon Sep 17 00:00:00 2001
From: Xiao Chen <xiao@apache.org>
Date: Fri, 18 May 2018 11:28:48 -0700
Subject: [PATCH] HADOOP-15154. Abstract new method assertCapability for
 StreamCapabilities testing. Contributed by Zsolt Venczel.

---
 .../hadoop/crypto/TestCryptoStreams.java      | 39 ++++++++++++-------
 .../hadoop/fs/contract/ContractTestUtils.java | 29 ++++++++++++++
 2 files changed, 54 insertions(+), 14 deletions(-)

diff --git a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/crypto/TestCryptoStreams.java b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/crypto/TestCryptoStreams.java
index 2172d8a3c55..cd7391a02c3 100644
--- a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/crypto/TestCryptoStreams.java
+++ b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/crypto/TestCryptoStreams.java
@@ -44,8 +44,7 @@ import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Test;
 
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
+import static org.apache.hadoop.fs.contract.ContractTestUtils.assertCapabilities;
 
 public class TestCryptoStreams extends CryptoStreamsTestBase {
   /**
@@ -419,21 +418,33 @@ public class TestCryptoStreams extends CryptoStreamsTestBase {
     // verify hasCapability returns what FakeOutputStream is set up for
     CryptoOutputStream cos =
         (CryptoOutputStream) getOutputStream(defaultBufferSize, key, iv);
-    assertTrue(cos instanceof StreamCapabilities);
-    assertTrue(cos.hasCapability(StreamCapabilities.HFLUSH));
-    assertTrue(cos.hasCapability(StreamCapabilities.HSYNC));
-    assertTrue(cos.hasCapability(StreamCapabilities.DROPBEHIND));
-    assertFalse(cos.hasCapability(StreamCapabilities.READAHEAD));
-    assertFalse(cos.hasCapability(StreamCapabilities.UNBUFFER));
+
+    assertCapabilities(cos,
+        new String[] {
+            StreamCapabilities.HFLUSH,
+            StreamCapabilities.HSYNC,
+            StreamCapabilities.DROPBEHIND
+        },
+        new String[] {
+            StreamCapabilities.READAHEAD,
+            StreamCapabilities.UNBUFFER
+        }
+    );
 
     // verify hasCapability for input stream
     CryptoInputStream cis =
         (CryptoInputStream) getInputStream(defaultBufferSize, key, iv);
-    assertTrue(cis instanceof StreamCapabilities);
-    assertTrue(cis.hasCapability(StreamCapabilities.DROPBEHIND));
-    assertTrue(cis.hasCapability(StreamCapabilities.READAHEAD));
-    assertTrue(cis.hasCapability(StreamCapabilities.UNBUFFER));
-    assertFalse(cis.hasCapability(StreamCapabilities.HFLUSH));
-    assertFalse(cis.hasCapability(StreamCapabilities.HSYNC));
+
+    assertCapabilities(cis,
+        new String[] {
+            StreamCapabilities.DROPBEHIND,
+            StreamCapabilities.READAHEAD,
+            StreamCapabilities.UNBUFFER
+        },
+        new String[] {
+            StreamCapabilities.HFLUSH,
+            StreamCapabilities.HSYNC
+        }
+    );
   }
 }
diff --git a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/fs/contract/ContractTestUtils.java b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/fs/contract/ContractTestUtils.java
index 54d015a9554..38a6fb10138 100644
--- a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/fs/contract/ContractTestUtils.java
+++ b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/fs/contract/ContractTestUtils.java
@@ -26,6 +26,7 @@ import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.LocatedFileStatus;
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.fs.RemoteIterator;
+import org.apache.hadoop.fs.StreamCapabilities;
 import org.apache.hadoop.io.IOUtils;
 import org.junit.Assert;
 import org.junit.internal.AssumptionViolatedException;
@@ -1437,6 +1438,34 @@ public class ContractTestUtils extends Assert {
     return list;
   }
 
+  /**
+   * Custom assert to test {@link StreamCapabilities}.
+   *
+   * @param stream The stream to test for StreamCapabilities
+   * @param shouldHaveCapabilities The array of expected capabilities
+   * @param shouldNotHaveCapabilities The array of unexpected capabilities
+   */
+  public static void assertCapabilities(
+      Object stream, String[] shouldHaveCapabilities,
+      String[] shouldNotHaveCapabilities) {
+    assertTrue("Stream should be instanceof StreamCapabilities",
+        stream instanceof StreamCapabilities);
+
+    if (shouldHaveCapabilities!=null) {
+      for (String shouldHaveCapability : shouldHaveCapabilities) {
+        assertTrue("Should have capability: " + shouldHaveCapability,
+            ((StreamCapabilities) stream).hasCapability(shouldHaveCapability));
+      }
+    }
+
+    if (shouldNotHaveCapabilities!=null) {
+      for (String shouldNotHaveCapability : shouldNotHaveCapabilities) {
+        assertFalse("Should not have capability: " + shouldNotHaveCapability,
+            ((StreamCapabilities) stream)
+                .hasCapability(shouldNotHaveCapability));
+      }
+    }
+  }
 
   /**
    * Results of recursive directory creation/scan operations.