From 168f3d39527e78aac800e89ce1397f9aa28e1000 Mon Sep 17 00:00:00 2001 From: Jay Patel Date: Fri, 7 Jun 2024 14:05:26 -0700 Subject: [PATCH] Fix tests and log requestId --- .../ingest/connection/RequestBuilder.java | 19 ++++++++++++------- .../SnowflakeStreamingIngestChannelTest.java | 17 +++++++++-------- .../SnowflakeStreamingIngestClientTest.java | 12 ++++++++---- 3 files changed, 29 insertions(+), 19 deletions(-) diff --git a/src/main/java/net/snowflake/ingest/connection/RequestBuilder.java b/src/main/java/net/snowflake/ingest/connection/RequestBuilder.java index 2e8b13553..fa2ac3ba7 100644 --- a/src/main/java/net/snowflake/ingest/connection/RequestBuilder.java +++ b/src/main/java/net/snowflake/ingest/connection/RequestBuilder.java @@ -678,18 +678,23 @@ public HttpGet generateHistoryRangeRequest( */ public HttpPost generateStreamingIngestPostRequest( String payload, String endPoint, String message) { - LOGGER.debug("Generate Snowpipe streaming request: endpoint={}, payload={}", endPoint, payload); + final String requestId = UUID.randomUUID().toString(); + LOGGER.debug( + "Generate Snowpipe streaming request: endpoint={}, payload={}, requestId={}", + endPoint, + payload, + requestId); // Make the corresponding URI URI uri = null; try { uri = new URIBuilder() - .setScheme(scheme) - .setHost(host) - .setPort(port) - .setPath(endPoint) - .setParameter(REQUEST_ID, UUID.randomUUID().toString()) - .build(); + .setScheme(scheme) + .setHost(host) + .setPort(port) + .setPath(endPoint) + .setParameter(REQUEST_ID, requestId) + .build(); } catch (URISyntaxException e) { throw new SFException(e, ErrorCode.BUILD_REQUEST_FAILURE, message); } diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java index 87e3f8f11..f025fc89f 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java @@ -1,12 +1,7 @@ package net.snowflake.ingest.streaming.internal; import static java.time.ZoneOffset.UTC; -import static net.snowflake.ingest.utils.Constants.ACCOUNT_URL; -import static net.snowflake.ingest.utils.Constants.OPEN_CHANNEL_ENDPOINT; -import static net.snowflake.ingest.utils.Constants.PRIVATE_KEY; -import static net.snowflake.ingest.utils.Constants.RESPONSE_SUCCESS; -import static net.snowflake.ingest.utils.Constants.ROLE; -import static net.snowflake.ingest.utils.Constants.USER; +import static net.snowflake.ingest.utils.Constants.*; import static org.mockito.ArgumentMatchers.argThat; import java.security.KeyPair; @@ -313,8 +308,14 @@ public void testOpenChannelPostRequest() throws Exception { requestBuilder.generateStreamingIngestPostRequest( payload, OPEN_CHANNEL_ENDPOINT, "open channel"); - Assert.assertEquals( - String.format("%s%s", urlStr, OPEN_CHANNEL_ENDPOINT), request.getRequestLine().getUri()); + String expectedUrlPattern = + String.format("%s%s", urlStr, OPEN_CHANNEL_ENDPOINT) + "(\\?requestId=[a-f0-9\\-]{36})?"; + + Assert.assertTrue( + String.format( + "Expected URL to match pattern: %s but was: %s", + expectedUrlPattern, request.getRequestLine().getUri()), + request.getRequestLine().getUri().matches(expectedUrlPattern)); Assert.assertNotNull(request.getFirstHeader(HttpHeaders.USER_AGENT)); Assert.assertNotNull(request.getFirstHeader(HttpHeaders.AUTHORIZATION)); Assert.assertEquals("POST", request.getMethod()); diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java index 1693e1520..5b24bcc7f 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java @@ -368,7 +368,6 @@ public void testGetChannelsStatusWithRequest() throws Exception { ChannelsStatusRequest.ChannelStatusRequestDTO dto = new ChannelsStatusRequest.ChannelStatusRequestDTO(channel); ChannelsStatusRequest request = new ChannelsStatusRequest(); - request.setRequestId("null_0"); request.setChannels(Collections.singletonList(dto)); ChannelsStatusResponse result = client.getChannelsStatus(Collections.singletonList(channel)); Assert.assertEquals(response.getMessage(), result.getMessage()); @@ -547,9 +546,15 @@ public void testRegisterBlobRequestCreationSuccess() throws Exception { HttpPost request = requestBuilder.generateStreamingIngestPostRequest( payload, REGISTER_BLOB_ENDPOINT, "register blob"); + String expectedUrlPattern = + String.format("%s%s", urlStr, REGISTER_BLOB_ENDPOINT) + "(\\?requestId=[a-f0-9\\-]{36})?"; + + Assert.assertTrue( + String.format( + "Expected URL to match pattern: %s but was: %s", + expectedUrlPattern, request.getRequestLine().getUri()), + request.getRequestLine().getUri().matches(expectedUrlPattern)); - Assert.assertEquals( - String.format("%s%s", urlStr, REGISTER_BLOB_ENDPOINT), request.getRequestLine().getUri()); Assert.assertNotNull(request.getFirstHeader(HttpHeaders.USER_AGENT)); Assert.assertNotNull(request.getFirstHeader(HttpHeaders.AUTHORIZATION)); Assert.assertEquals("POST", request.getMethod()); @@ -1432,7 +1437,6 @@ public void testGetLatestCommittedOffsetTokens() throws Exception { ChannelsStatusRequest.ChannelStatusRequestDTO dto = new ChannelsStatusRequest.ChannelStatusRequestDTO(channel); ChannelsStatusRequest request = new ChannelsStatusRequest(); - request.setRequestId("null_0"); request.setChannels(Collections.singletonList(dto)); Map result = client.getLatestCommittedOffsetTokens(Collections.singletonList(channel));