Stabilize most commonly failing tests (Azure#36956)

Stabilize most commonly failing tests

Author: alzimmermsft

sdk/attestation/azure-security-attestation/src/test/java/com/azure/security/attestation/models/AttestationTokenTests.java

@@ -14,6 +14,8 @@
 import java.security.cert.X509Certificate;
 import java.time.OffsetDateTime;
 import java.time.ZoneOffset;
+import java.time.format.DateTimeFormatter;
+import java.time.temporal.ChronoUnit;
 import java.util.LinkedHashMap;
 
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
@@ -24,6 +26,7 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
 
 /**
  * Test for Attestation Signing Certificates APIs.
@@ -269,8 +272,7 @@ void verifyAttestationTokenIssuer() {
 
     @Test
     void verifyAttestationTokenExpireTimeout() {
-        OffsetDateTime timeNow = OffsetDateTime.now();
-        timeNow = timeNow.minusNanos(timeNow.getNano());
+        final OffsetDateTime timeNow = OffsetDateTime.now().truncatedTo(ChronoUnit.MICROS);
 
         TestObject testObjectExpired30SecondsAgo = new TestObject()
             .setAlg("Test Algorithm")
@@ -290,14 +292,36 @@ void verifyAttestationTokenExpireTimeout() {
             ((AttestationTokenImpl) newToken).validate(null,
                 new AttestationTokenValidationOptions()));
         // Both the current time and the expiration time should be in the exception message.
-        assertTrue(ex.getMessage().contains("expiration"));
+        String exceptionMessage = ex.getMessage();
+        assertTrue(exceptionMessage.contains("expiration"), () ->
+            "Expected exception message to contain 'expiration' but it didn't. Actual exception message: "
+            + exceptionMessage);
         // Because the TestObject round-trips times through Epoch times, they are in UTC time.
         // Adjust the target time to be in UTC rather than the current time zone, since we're checking to ensure
         // that the time is reflected in the exception message.
         OffsetDateTime expTime = timeNow.minusSeconds(30).withOffsetSameInstant(ZoneOffset.UTC);
 
-        assertTrue(ex.getMessage().contains(String.format("%tc", timeNow)));
-        assertTrue(ex.getMessage().contains(String.format("%tc", expTime)));
+        // Format of the exception message is "Current time: <current time> Expiration time: <expiration time>"
+        // Since the test could take a while and the current time is based on the time when the exception is thrown
+        // this can cause it to be different than 'timeNow' when the test started.
+        // To make sure this test isn't flaky capture the datetime string in the exception message, turn it into an
+        // OffsetDateTime and compare it to 'timeNow' allowing for some skew.
+        // Date format is 'Wed Sep 27 12:48:15 -04:00 2023'
+        DateTimeFormatter formatter = DateTimeFormatter.ofPattern("EEE MMM dd HH:mm:ss XXX yyyy");
+
+        int currentTimeIndex = exceptionMessage.indexOf("Current time: ");
+        int expirationTimeIndex = exceptionMessage.indexOf("Expiration time: ");
+        String currentTimeInExceptionString = exceptionMessage.substring(currentTimeIndex + 14, expirationTimeIndex - 1);
+        OffsetDateTime currentTimeInException = OffsetDateTime.parse(currentTimeInExceptionString, formatter);
+        long skew = timeNow.until(currentTimeInException, ChronoUnit.SECONDS);
+        if (skew > 5 || skew < 0) {
+            fail(String.format("Expected exception message to contain 'Current Time' within 5 seconds, but not before, "
+                               + "of %tc but it was greater. Actual exception message: %s", timeNow, exceptionMessage));
+        }
+
+        assertTrue(exceptionMessage.contains(String.format("%tc", expTime)), () -> String.format(
+            "Expected exception message to contain '%tc' but it didn't. Actual exception message: %s", expTime,
+            exceptionMessage));
     }
 
     @Test

sdk/core/azure-core-http-netty/src/main/java/com/azure/core/http/netty/NettyAsyncHttpClient.java

@@ -26,7 +26,9 @@
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
 import io.netty.channel.EventLoopGroup;
+import io.netty.handler.codec.http.HttpHeaders;
 import io.netty.handler.codec.http.HttpMethod;
+import io.netty.handler.proxy.HttpProxyHandler;
 import io.netty.handler.proxy.ProxyConnectException;
 import io.netty.handler.stream.ChunkedNioFile;
 import io.netty.handler.stream.ChunkedStream;
@@ -40,6 +42,8 @@
 import reactor.netty.NettyPipeline;
 import reactor.netty.http.client.HttpClientRequest;
 import reactor.netty.http.client.HttpClientResponse;
+import reactor.util.function.Tuple2;
+import reactor.util.function.Tuples;
 
 import javax.net.ssl.SSLException;
 import java.io.IOException;
@@ -120,7 +124,8 @@ public Mono<HttpResponse> send(HttpRequest request, Context context) {
     private Mono<HttpResponse> attemptAsync(HttpRequest request, boolean eagerlyReadResponse,
         boolean ignoreResponseBody, boolean headersEagerlyConverted, Long responseTimeout,
         ProgressReporter progressReporter, boolean proxyRetry) {
-        Flux<HttpResponse> nettyRequest = nettyClient.request(toReactorNettyHttpMethod(request.getHttpMethod()))
+        Flux<Tuple2<HttpResponse, HttpHeaders>> nettyRequest = nettyClient
+            .request(toReactorNettyHttpMethod(request.getHttpMethod()))
             .uri(request.getUrl().toString())
             .send(bodySendDelegate(request))
             .responseConnection(responseDelegate(request, disableBufferCopy, eagerlyReadResponse, ignoreResponseBody,
@@ -132,11 +137,18 @@ private Mono<HttpResponse> attemptAsync(HttpRequest request, boolean eagerlyRead
         }
 
         return nettyRequest.single()
-            .flatMap(response -> {
+            .flatMap(responseAndHeaders -> {
+                HttpResponse response = responseAndHeaders.getT1();
                 if (addProxyHandler && response.getStatusCode() == 407) {
-                    return proxyRetry
-                        ? Mono.error(new ProxyConnectException("Connection to proxy failed."))
-                        : Mono.error(new ProxyConnectException("First attempt to connect to proxy failed."));
+                    if (proxyRetry) {
+                        // Exhausted retry attempt return an error.
+                        return Mono.error(new HttpProxyHandler.HttpProxyConnectException(
+                            "Failed to connect to proxy. Status: 407", responseAndHeaders.getT2()));
+                    } else {
+                        // Retry the request.
+                        return attemptAsync(request, eagerlyReadResponse, ignoreResponseBody, headersEagerlyConverted,
+                            responseTimeout, progressReporter, true);
+                    }
                 } else {
                     return Mono.just(response);
                 }
@@ -285,7 +297,7 @@ private static NettyOutbound sendInputStream(NettyOutbound reactorNettyOutbound,
      * HttpHeaders.
      * @return a delegate upon invocation setup Rest response object
      */
-    private static BiFunction<HttpClientResponse, Connection, Mono<HttpResponse>> responseDelegate(
+    private static BiFunction<HttpClientResponse, Connection, Mono<Tuple2<HttpResponse, HttpHeaders>>> responseDelegate(
         HttpRequest restRequest, boolean disableBufferCopy, boolean eagerlyReadResponse, boolean ignoreResponseBody,
         boolean headersEagerlyConverted) {
         return (reactorNettyResponse, reactorNettyConnection) -> {
@@ -314,11 +326,11 @@ private static BiFunction<HttpClientResponse, Connection, Mono<HttpResponse>> re
                 return reactorNettyConnection.inbound().receive().aggregate().asByteArray()
                     .doFinally(ignored -> closeConnection(reactorNettyConnection))
                     .switchIfEmpty(Mono.just(EMPTY_BYTES))
-                    .map(bytes -> new NettyAsyncHttpBufferedResponse(reactorNettyResponse, restRequest, bytes,
-                        headersEagerlyConverted));
+                    .map(bytes -> Tuples.of(new NettyAsyncHttpBufferedResponse(reactorNettyResponse, restRequest, bytes,
+                        headersEagerlyConverted), reactorNettyResponse.responseHeaders()));
             } else {
-                return Mono.just(new NettyAsyncHttpResponse(reactorNettyResponse, reactorNettyConnection, restRequest,
-                    disableBufferCopy, headersEagerlyConverted));
+                return Mono.just(Tuples.of(new NettyAsyncHttpResponse(reactorNettyResponse, reactorNettyConnection,
+                    restRequest, disableBufferCopy, headersEagerlyConverted), reactorNettyResponse.responseHeaders()));
             }
         };
     }

sdk/core/azure-core-http-netty/src/main/java/com/azure/core/http/netty/implementation/HttpProxyHandler.java

@@ -45,7 +45,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 
@@ -94,7 +93,6 @@ public final class HttpProxyHandler extends ProxyHandler {
     private final AuthorizationChallengeHandler challengeHandler;
     private final AtomicReference<ChallengeHolder> proxyChallengeHolderReference;
     private final HttpClientCodec codec;
-    private final AtomicBoolean hasHandledChallenge = new AtomicBoolean();
 
     private String authScheme = null;
     private HttpResponseStatus status;
@@ -190,7 +188,6 @@ private String createAuthorizationHeader() {
          * created from the same client.
          */
         if (proxyChallengeHolder != null) {
-            hasHandledChallenge.set(true);
             // Attempt to apply digest challenges, these are preferred over basic authorization.
             List<Map<String, String>> digestChallenges = proxyChallengeHolder.getDigestChallenges();
             if (!CoreUtils.isNullOrEmpty(digestChallenges)) {
@@ -238,19 +235,9 @@ protected boolean handleResponse(ChannelHandlerContext ctx, Object o) throws Pro
         }
 
         boolean responseComplete = o instanceof LastHttpContent;
-        if (responseComplete) {
-            if (status == null) {
-                throw new io.netty.handler.proxy.HttpProxyHandler.HttpProxyConnectException(
-                    "Never received response for CONNECT request.", innerHeaders);
-            } else if (status.code() != 200) {
-                // Return the error response on the first attempt as the proxy handler doesn't apply credentials on the
-                // first attempt.
-                if (hasHandledChallenge.get()) {
-                    // Later attempts throw an exception.
-                    throw new io.netty.handler.proxy.HttpProxyHandler.HttpProxyConnectException(
-                        "Failed to connect to proxy. Status: " + status, innerHeaders);
-                }
-            }
+        if (responseComplete && status == null) {
+            throw new io.netty.handler.proxy.HttpProxyHandler.HttpProxyConnectException(
+                "Never received response for CONNECT request.", innerHeaders);
         }
 
         return responseComplete;

sdk/core/azure-core/src/main/java/com/azure/core/http/ProxyOptions.java

@@ -22,7 +22,7 @@
 import java.util.regex.PatternSyntaxException;
 
 /**
- * This represents proxy configuration to be used in http clients..
+ * This represents proxy configuration to be used in http clients.
  */
 public class ProxyOptions {
     private static final ClientLogger LOGGER = new ClientLogger(ProxyOptions.class);
@@ -136,7 +136,7 @@ public InetSocketAddress getAddress() {
     }
 
     /**
-     * Gets the type of the prxoy.
+     * Gets the type of the proxy.
      *
      * @return the type of the proxy.
      */

sdk/core/azure-core/src/main/java/com/azure/core/http/policy/HttpLoggingPolicy.java

@@ -246,16 +246,14 @@ private void logBody(HttpRequest request, int contentLength, LoggingEventBuilder
         } else {
             // Add non-mutating operators to the data stream.
             AccessibleByteArrayOutputStream stream = new AccessibleByteArrayOutputStream(contentLength);
-            request.setBody(
-                content.toFluxByteBuffer()
-                    .doOnNext(byteBuffer -> {
-                        try {
-                            ImplUtils.writeByteBufferToStream(byteBuffer.duplicate(), stream);
-                        } catch (IOException ex) {
-                            throw LOGGER.logExceptionAsError(new UncheckedIOException(ex));
-                        }
-                    })
-                    .doFinally(ignored -> logBody(logBuilder, logger, contentType, stream.toString(StandardCharsets.UTF_8))));
+            request.setBody(Flux.using(() -> content, con -> con.toFluxByteBuffer()
+                .doOnNext(byteBuffer -> {
+                    try {
+                        ImplUtils.writeByteBufferToStream(byteBuffer.duplicate(), stream);
+                    } catch (IOException ex) {
+                        throw LOGGER.logExceptionAsError(new UncheckedIOException(ex));
+                    }
+                }), ignored -> logBody(logBuilder, logger, contentType, stream.toString(StandardCharsets.UTF_8))));
         }
     }
 
@@ -548,15 +546,14 @@ public HttpHeaders getHeaders() {
         public Flux<ByteBuffer> getBody() {
             AccessibleByteArrayOutputStream stream = new AccessibleByteArrayOutputStream(contentLength);
 
-            return actualResponse.getBody()
+            return Flux.using(() -> actualResponse, response -> response.getBody()
                 .doOnNext(byteBuffer -> {
                     try {
                         ImplUtils.writeByteBufferToStream(byteBuffer.duplicate(), stream);
                     } catch (IOException ex) {
                         throw LOGGER.logExceptionAsError(new UncheckedIOException(ex));
                     }
-                })
-                .doFinally(ignored -> doLog(stream.toString(StandardCharsets.UTF_8)));
+                }), ignored -> doLog(stream.toString(StandardCharsets.UTF_8)));
         }
 
         @Override

sdk/core/azure-core/src/main/java/com/azure/core/implementation/logging/DefaultLogger.java

@@ -12,16 +12,16 @@
 
 import java.io.PrintWriter;
 import java.io.StringWriter;
+import java.nio.charset.StandardCharsets;
 import java.nio.file.InvalidPathException;
 import java.time.LocalDateTime;
-import java.time.format.DateTimeFormatter;
+import java.time.temporal.ChronoField;
 
 /**
  * This class is an internal implementation of slf4j logger.
  */
 public final class DefaultLogger extends MarkerIgnoringBase {
     private static final long serialVersionUID = -144261058636441630L;
-    private static final DateTimeFormatter DATE_FORMAT = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSS");
 
     // The template for the log message:
     // YYYY-MM-DD HH:MM:ss.SSS [thread] [level] classpath - message
@@ -49,7 +49,7 @@ public final class DefaultLogger extends MarkerIgnoringBase {
      * @param clazz Class creating the logger.
      */
     public DefaultLogger(Class<?> clazz) {
-        this(clazz.getName());
+        this(clazz.getCanonicalName(), false);
     }
 
     /**
@@ -59,14 +59,20 @@ public DefaultLogger(Class<?> clazz) {
      * name passes in.
      */
     public DefaultLogger(String className) {
-        String classPath;
+        this(getClassPathFromClassName(className), false);
+    }
+
+    private static String getClassPathFromClassName(String className) {
         try {
-            classPath = Class.forName(className).getCanonicalName();
+            return Class.forName(className).getCanonicalName();
         } catch (ClassNotFoundException | InvalidPathException e) {
             // Swallow ClassNotFoundException as the passed class name may not correlate to an actual class.
             // Swallow InvalidPathException as the className may contain characters that aren't legal file characters.
-            classPath = className;
+            return className;
         }
+    }
+
+    private DefaultLogger(String classPath, boolean ignored) {
         this.classPath = classPath;
         int configuredLogLevel = fromEnvironment().getLogLevel();
 
@@ -75,7 +81,6 @@ public DefaultLogger(String className) {
         isInfoEnabled = LogLevel.INFORMATIONAL.getLogLevel() >= configuredLogLevel;
         isWarnEnabled = LogLevel.WARNING.getLogLevel() >= configuredLogLevel;
         isErrorEnabled = LogLevel.ERROR.getLogLevel() >= configuredLogLevel;
-
     }
 
     private static LogLevel fromEnvironment() {
@@ -377,9 +382,57 @@ private void log(String levelName, String message, Throwable t) {
      *
      * @return The current time in {@code DATE_FORMAT}
      */
-    private String getFormattedDate() {
+    private static String getFormattedDate() {
         LocalDateTime now = LocalDateTime.now();
-        return DATE_FORMAT.format(now);
+
+        // yyyy-MM-dd HH:mm:ss.SSS
+        // 23 characters that will be ASCII
+        byte[] bytes = new byte[23];
+
+        // yyyy-
+        int year = now.getYear();
+        int round = year / 1000;
+        bytes[0] = (byte) ('0' + round);
+        year = year - (1000 * round);
+        round = year / 100;
+        bytes[1] = (byte) ('0' + round);
+        year = year - (100 * round);
+        round = year / 10;
+        bytes[2] = (byte) ('0' + round);
+        bytes[3] = (byte) ('0' + (year - (10 * round)));
+        bytes[4] = '-';
+
+        // MM-
+        zeroPad(now.getDayOfMonth(), bytes, 5);
+        bytes[7] = '-';
+
+        // dd
+        zeroPad(now.getDayOfMonth(), bytes, 8);
+        bytes[10] = ' ';
+
+        // HH:
+        zeroPad(now.getHour(), bytes, 11);
+        bytes[13] = ':';
+
+        // mm:
+        zeroPad(now.getMinute(), bytes, 14);
+        bytes[16] = ':';
+
+        // ss.
+        zeroPad(now.getSecond(), bytes, 17);
+        bytes[19] = '.';
+
+        // SSS
+        int millis = now.get(ChronoField.MILLI_OF_SECOND);
+        round = millis / 100;
+        bytes[20] = (byte) ('0' + round);
+        millis = millis - (100 * round);
+        round = millis / 10;
+        bytes[21] = (byte) ('0' + round);
+        bytes[22] = (byte) ('0' + (millis - (10 * round)));
+
+        // Use UTF-8 as it's more performant than ASCII in Java 8
+        return new String(bytes, StandardCharsets.UTF_8);
     }
 
     /**
@@ -398,4 +451,15 @@ void writeWithThrowable(StringBuilder stringBuilder, Throwable t) {
         }
         System.out.print(stringBuilder.toString());
     }
+
+    private static void zeroPad(int value, byte[] bytes, int index) {
+        if (value < 10) {
+            bytes[index++] = '0';
+            bytes[index] = (byte) ('0' + value);
+        } else {
+            int high = value / 10;
+            bytes[index++] = (byte) ('0' + high);
+            bytes[index] = (byte) ('0' + (value - (10 * high)));
+        }
+    }
 }