Fix Exception in Test Recording (#24675)

* Fix Exception in Test Recording * Roll-back version change
Azure · Oct 8, 2021 · 297922c · 297922c
1 parent cd25021
commit 297922c
Show file tree

Hide file tree

Showing 6 changed files with 19 additions and 16 deletions.
diff --git a/...ore/azure-core-test/src/main/java/com/azure/core/test/policy/RecordNetworkCallPolicy.java b/...ore/azure-core-test/src/main/java/com/azure/core/test/policy/RecordNetworkCallPolicy.java
@@ -17,6 +17,7 @@
 import com.azure.core.test.models.RecordedData;
 import com.azure.core.test.models.RecordingRedactor;
 import com.azure.core.util.CoreUtils;
+import com.azure.core.util.FluxUtil;
 import com.azure.core.util.UrlBuilder;
 import com.azure.core.util.logging.ClientLogger;
 import reactor.core.Exceptions;
@@ -180,9 +181,9 @@ private static Mono<Tuple2<HttpResponse, Map<String, String>>> extractResponseDa
         }
 
         final HttpResponse bufferedResponse = response.buffer();
+        final Mono<byte[]> responseBody = FluxUtil.collectBytesInByteBufferStream(bufferedResponse.getBody());
         if (contentType == null) {
-            return bufferedResponse.getBodyAsByteArray()
-                .switchIfEmpty(Mono.defer(() -> Mono.just(new byte[0])))
+            return responseBody.switchIfEmpty(Mono.defer(() -> Mono.just(new byte[0])))
                 .map(bytes -> {
                     if (bytes.length == 0) {
                         return Tuples.of(bufferedResponse, responseData);
@@ -195,8 +196,7 @@ private static Mono<Tuple2<HttpResponse, Map<String, String>>> extractResponseDa
                 });
         } else if (contentType.equalsIgnoreCase(ContentType.APPLICATION_OCTET_STREAM)
             || contentType.equalsIgnoreCase("avro/binary")) {
-            return bufferedResponse.getBodyAsByteArray()
-                .switchIfEmpty(Mono.defer(() -> Mono.just(new byte[0])))
+            return responseBody.switchIfEmpty(Mono.defer(() -> Mono.just(new byte[0])))
                 .map(bytes -> {
                     if (bytes.length == 0) {
                         return Tuples.of(bufferedResponse, responseData);
@@ -206,15 +206,14 @@ private static Mono<Tuple2<HttpResponse, Map<String, String>>> extractResponseDa
                     return Tuples.of(bufferedResponse, responseData);
                 });
         } else if (contentType.contains("json") || response.getHeaderValue(CONTENT_ENCODING) == null) {
-            return bufferedResponse.getBodyAsString(StandardCharsets.UTF_8)
+            return responseBody.map(bytes -> CoreUtils.bomAwareToString(bytes, response.getHeaderValue(CONTENT_TYPE)))
                 .switchIfEmpty(Mono.defer(() -> Mono.just("")))
                 .map(content -> {
                     responseData.put(BODY, redactor.redact(content));
                     return Tuples.of(bufferedResponse, responseData);
                 });
         } else {
-            return bufferedResponse.getBodyAsByteArray()
-                .switchIfEmpty(Mono.defer(() -> Mono.just(new byte[0])))
+            return responseBody.switchIfEmpty(Mono.defer(() -> Mono.just(new byte[0])))
                 .map(bytes -> {
                     if (bytes.length == 0) {
                         return Tuples.of(bufferedResponse, responseData);

diff --git a/sdk/tables/azure-data-tables/src/test/java/com/azure/data/tables/TableAsyncClientTest.java b/sdk/tables/azure-data-tables/src/test/java/com/azure/data/tables/TableAsyncClientTest.java
@@ -55,8 +55,7 @@
 public class TableAsyncClientTest extends TestBase {
     private static final Duration TIMEOUT = Duration.ofSeconds(100);
     private static final HttpClient DEFAULT_HTTP_CLIENT = HttpClient.createDefault();
-    private static final boolean IS_COSMOS_TEST = System.getenv("AZURE_TABLES_CONNECTION_STRING") != null
-        && System.getenv("AZURE_TABLES_CONNECTION_STRING").contains("cosmos.azure.com");
+    private static final boolean IS_COSMOS_TEST = TestUtils.isCosmosTest();
 
     private TableAsyncClient tableClient;
     private HttpPipelinePolicy recordPolicy;

diff --git a/sdk/tables/azure-data-tables/src/test/java/com/azure/data/tables/TableClientTest.java b/sdk/tables/azure-data-tables/src/test/java/com/azure/data/tables/TableClientTest.java
@@ -54,8 +54,8 @@
  */
 public class TableClientTest extends TestBase {
     private static final HttpClient DEFAULT_HTTP_CLIENT = HttpClient.createDefault();
-    private static final boolean IS_COSMOS_TEST = System.getenv("AZURE_TABLES_CONNECTION_STRING") != null
-        && System.getenv("AZURE_TABLES_CONNECTION_STRING").contains("cosmos.azure.com");
+    private static final boolean IS_COSMOS_TEST = TestUtils.isCosmosTest();
+
 
     private TableClient tableClient;
     private HttpPipelinePolicy recordPolicy;

diff --git a/...es/azure-data-tables/src/test/java/com/azure/data/tables/TableServiceAsyncClientTest.java b/...es/azure-data-tables/src/test/java/com/azure/data/tables/TableServiceAsyncClientTest.java
@@ -52,8 +52,7 @@
 public class TableServiceAsyncClientTest extends TestBase {
     private static final Duration TIMEOUT = Duration.ofSeconds(100);
     private static final HttpClient DEFAULT_HTTP_CLIENT = HttpClient.createDefault();
-    private static final boolean IS_COSMOS_TEST = System.getenv("AZURE_TABLES_CONNECTION_STRING") != null
-        && System.getenv("AZURE_TABLES_CONNECTION_STRING").contains("cosmos.azure.com");
+    private static final boolean IS_COSMOS_TEST = TestUtils.isCosmosTest();
 
     private TableServiceAsyncClient serviceClient;
     private HttpPipelinePolicy recordPolicy;

diff --git a/sdk/tables/azure-data-tables/src/test/java/com/azure/data/tables/TableServiceClientTest.java b/sdk/tables/azure-data-tables/src/test/java/com/azure/data/tables/TableServiceClientTest.java
@@ -55,8 +55,7 @@
  */
 public class TableServiceClientTest extends TestBase {
     private static final HttpClient DEFAULT_HTTP_CLIENT = HttpClient.createDefault();
-    private static final boolean IS_COSMOS_TEST = System.getenv("AZURE_TABLES_CONNECTION_STRING") != null
-        && System.getenv("AZURE_TABLES_CONNECTION_STRING").contains("cosmos.azure.com");
+    private static final boolean IS_COSMOS_TEST = TestUtils.isCosmosTest();
 
     private TableServiceClient serviceClient;
     private HttpPipelinePolicy recordPolicy;

diff --git a/sdk/tables/azure-data-tables/src/test/java/com/azure/data/tables/TestUtils.java b/sdk/tables/azure-data-tables/src/test/java/com/azure/data/tables/TestUtils.java
@@ -13,6 +13,7 @@
 import com.azure.core.http.HttpResponse;
 import com.azure.core.http.policy.HttpPipelinePolicy;
 import com.azure.core.test.http.MockHttpResponse;
+import com.azure.core.util.Configuration;
 import com.azure.core.util.CoreUtils;
 import com.azure.core.util.DateTimeRfc1123;
 import com.azure.data.tables.models.TableServiceProperties;
@@ -46,7 +47,7 @@ private TestUtils() {
     public static String getConnectionString(boolean isPlaybackMode) {
         return isPlaybackMode
             ? "DefaultEndpointsProtocol=https;AccountName=dummyAccount;AccountKey=xyzDummy;EndpointSuffix=core.windows.net"
-            : System.getenv("AZURE_TABLES_CONNECTION_STRING");
+            : Configuration.getGlobalConfiguration().get("AZURE_TABLES_CONNECTION_STRING");
     }
 
     public static HttpRequest request(String url) throws MalformedURLException {
@@ -185,4 +186,10 @@ static void assertPropertiesEquals(TableServiceProperties expected,
             assertNull(actual.getMinuteMetrics());
         }
     }
+
+    static boolean isCosmosTest() {
+        Configuration globalConfiguration = Configuration.getGlobalConfiguration();
+        return globalConfiguration.get("AZURE_TABLES_CONNECTION_STRING") != null
+            && globalConfiguration.get("AZURE_TABLES_CONNECTION_STRING").contains("cosmos.azure.com");
+    }
 }