Skip to content

Commit b0aa824

Browse files
Trianz-AkshayTrianz-Akshay
committed
add unit test cases
1 parent dace55b commit b0aa824

File tree

3 files changed

+239
-0
lines changed

3 files changed

+239
-0
lines changed

‎athena-datalakegen2/src/test/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2MetadataHandlerTest.java‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@
7070

7171
import static com.amazonaws.athena.connectors.datalakegen2.DataLakeGen2MetadataHandler.PARTITION_NUMBER;
7272
import static org.junit.Assert.assertEquals;
73+
import static org.junit.Assert.assertNotNull;
7374
import static org.mockito.ArgumentMatchers.any;
7475
import static org.mockito.ArgumentMatchers.nullable;
7576
import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
@@ -233,6 +234,88 @@ public void doGetTable()
233234
assertEquals("testCatalog", getTableResponse.getCatalogName());
234235
}
235236

237+
@Test
238+
public void testGetSchemaWithAzureServerlessEnvironment()
239+
throws Exception
240+
{
241+
BlockAllocator blockAllocator = new BlockAllocatorImpl();
242+
243+
// Mock the data type query result set for Azure serverless
244+
String[] dataTypeSchema = {"COLUMN_NAME", "DATA_TYPE"};
245+
Object[][] dataTypeValues = {
246+
{"id", "int"},
247+
{"name", "varchar"},
248+
{"email", "nvarchar"},
249+
{"created_date", "datetime"},
250+
{"is_active", "bit"},
251+
{"salary", "decimal"},
252+
{"binary_data", "varbinary"}
253+
};
254+
AtomicInteger dataTypeRowNumber = new AtomicInteger(-1);
255+
ResultSet dataTypeResultSet = mockResultSet(dataTypeSchema, dataTypeValues, dataTypeRowNumber);
256+
257+
// Mock the prepared statement and its execution
258+
java.sql.PreparedStatement mockPreparedStatement = mock(java.sql.PreparedStatement.class);
259+
when(connection.prepareStatement(any(String.class))).thenReturn(mockPreparedStatement);
260+
when(mockPreparedStatement.executeQuery()).thenReturn(dataTypeResultSet);
261+
262+
// Mock Azure serverless URL
263+
when(connection.getMetaData().getURL()).thenReturn("jdbc:sqlserver://myworkspace-ondemand.sql.azuresynapse.net:1433;database=mydatabase;");
264+
when(connection.getCatalog()).thenReturn("testCatalog");
265+
266+
TableName inputTableName = new TableName("TESTSCHEMA", "TESTTABLE");
267+
GetTableResponse getTableResponse = this.dataLakeGen2MetadataHandler.doGetTable(
268+
blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap()));
269+
270+
// Verify the response
271+
assertNotNull(getTableResponse);
272+
assertEquals(inputTableName, getTableResponse.getTableName());
273+
assertEquals("testCatalog", getTableResponse.getCatalogName());
274+
275+
// Verify that the schema was built using the direct SQL query approach (Azure serverless path)
276+
Schema responseSchema = getTableResponse.getSchema();
277+
assertNotNull(responseSchema);
278+
assertEquals(8, responseSchema.getFields().size()); // 7 columns from our mock data + 1 partition field
279+
}
280+
281+
@Test
282+
public void testGetSchemaWithStandardEnvironment()
283+
throws Exception
284+
{
285+
BlockAllocator blockAllocator = new BlockAllocatorImpl();
286+
287+
// Mock the data type query result set
288+
String[] dataTypeSchema = {"COLUMN_NAME", "DATA_TYPE"};
289+
Object[][] dataTypeValues = {{"testCol1", "int"}, {"testCol2", "varchar"}};
290+
AtomicInteger dataTypeRowNumber = new AtomicInteger(-1);
291+
ResultSet dataTypeResultSet = mockResultSet(dataTypeSchema, dataTypeValues, dataTypeRowNumber);
292+
293+
// Mock the prepared statement and its execution
294+
java.sql.PreparedStatement mockPreparedStatement = mock(java.sql.PreparedStatement.class);
295+
when(connection.prepareStatement(any(String.class))).thenReturn(mockPreparedStatement);
296+
when(mockPreparedStatement.executeQuery()).thenReturn(dataTypeResultSet);
297+
298+
// Mock standard SQL Server URL (non-serverless)
299+
when(connection.getMetaData().getURL()).thenReturn("jdbc:sqlserver://myserver.database.windows.net:1433;database=mydatabase;");
300+
when(connection.getCatalog()).thenReturn("testCatalog");
301+
302+
// Mock the getColumns result set for standard environment
303+
String[] schema = {"TABLE_CAT", "TABLE_SCHEM", "TABLE_NAME", "COLUMN_NAME", "DATA_TYPE", "TYPE_NAME", "COLUMN_SIZE", "BUFFER_LENGTH", "DECIMAL_DIGITS", "NUM_PREC_RADIX", "NULLABLE", "REMARKS", "COLUMN_DEF", "SQL_DATA_TYPE", "SQL_DATETIME_SUB", "CHAR_OCTET_LENGTH", "ORDINAL_POSITION", "IS_NULLABLE"};
304+
Object[][] values = {{"testCatalog", "TESTSCHEMA", "TESTTABLE", "testCol1", Types.INTEGER, "int", 10, 4, 0, 10, 1, "", "", Types.INTEGER, 0, 4, 1, "YES"}};
305+
AtomicInteger rowNumber = new AtomicInteger(-1);
306+
ResultSet getColumnsResultSet = mockResultSet(schema, values, rowNumber);
307+
when(connection.getMetaData().getColumns("testCatalog", "TESTSCHEMA", "TESTTABLE", null)).thenReturn(getColumnsResultSet);
308+
309+
TableName inputTableName = new TableName("TESTSCHEMA", "TESTTABLE");
310+
GetTableResponse getTableResponse = this.dataLakeGen2MetadataHandler.doGetTable(
311+
blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName, Collections.emptyMap()));
312+
313+
// Verify the response
314+
assertNotNull(getTableResponse);
315+
assertEquals(inputTableName, getTableResponse.getTableName());
316+
assertEquals("testCatalog", getTableResponse.getCatalogName());
317+
}
318+
236319
@Test
237320
public void doListTables() throws Exception
238321
{

‎athena-datalakegen2/src/test/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2UtilTest.java‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
/*-
2+
* #%L
3+
* athena-datalakegen2
4+
* %%
5+
* Copyright (C) 2019 - 2025 Amazon Web Services
6+
* %%
7+
* Licensed under the Apache License, Version 2.0 (the "License");
8+
* you may not use this file except in compliance with the License.
9+
* You may obtain a copy of the License at
10+
*
11+
* http://www.apache.org/licenses/LICENSE-2.0
12+
*
13+
* Unless required by applicable law or agreed to in writing, software
14+
* distributed under the License is distributed on an "AS IS" BASIS,
15+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16+
* See the License for the specific language governing permissions and
17+
* limitations under the License.
18+
* #L%
19+
*/
20+
package com.amazonaws.athena.connectors.datalakegen2;
21+
22+
import org.junit.Test;
23+
24+
import static org.junit.Assert.assertEquals;
25+
import static org.junit.Assert.assertNull;
26+
27+
public class DataLakeGen2UtilTest
28+
{
29+
@Test
30+
public void testCheckEnvironmentWithAzureServerlessUrl()
31+
{
32+
// Test with Azure serverless URL containing "ondemand" in hostname
33+
String serverlessUrl = "datalakegentwo://jdbc:sqlserver://myworkspace-ondemand.sql.azuresynapse.net:1433;database=mydatabase;encrypt=true;trustServerCertificate=false;hostNameInCertificate=*.sql.azuresynapse.net;loginTimeout=30;";
34+
String result = DataLakeGen2Util.checkEnvironment(serverlessUrl);
35+
assertEquals("azureServerless", result);
36+
}
37+
38+
@Test
39+
public void testCheckEnvironmentWithStandardUrl()
40+
{
41+
// Test with standard SQL Server URL (should not detect as serverless)
42+
String standardUrl = "datalakegentwo://jdbc:sqlserver://myserver.database.windows.net:1433;database=mydatabase;encrypt=true;trustServerCertificate=false;hostNameInCertificate=*.database.windows.net;loginTimeout=30;";
43+
String result = DataLakeGen2Util.checkEnvironment(standardUrl);
44+
assertNull(result);
45+
}
46+
47+
@Test
48+
public void testCheckEnvironmentWithOnDemandInDatabaseName()
49+
{
50+
// Test with "ondemand" in database name but not hostname (should not detect as serverless)
51+
String urlWithOnDemandInDb = "datalakegentwo://jdbc:sqlserver://myserver.database.windows.net:1433;database=ondemand_database;";
52+
String result = DataLakeGen2Util.checkEnvironment(urlWithOnDemandInDb);
53+
assertNull(result);
54+
}
55+
56+
@Test
57+
public void testCheckEnvironmentWithNullAndEmptyUrls()
58+
{
59+
// Test with null, empty, and invalid URLs
60+
assertNull(DataLakeGen2Util.checkEnvironment(null));
61+
assertNull(DataLakeGen2Util.checkEnvironment(""));
62+
assertNull(DataLakeGen2Util.checkEnvironment(" "));
63+
assertNull(DataLakeGen2Util.checkEnvironment("invalid-url-format"));
64+
assertNull(DataLakeGen2Util.checkEnvironment("https://ONDEMAND-server.database.windows.net:1433;database=mydatabase;"));
65+
}
66+
}

‎athena-datalakegen2/src/test/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeRecordHandlerTest.java‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient;
4444

4545
import java.sql.Connection;
46+
import java.sql.DatabaseMetaData;
4647
import java.sql.PreparedStatement;
4748
import java.sql.SQLException;
4849
import java.util.Collections;
@@ -119,4 +120,93 @@ public void buildSplitSqlNew()
119120
Mockito.verify(preparedStatement, Mockito.times(1)).setString(1, "varcharTest");
120121
}
121122

123+
@Test
124+
public void testReadWithConstraintWithAzureServerlessEnvironment()
125+
throws Exception
126+
{
127+
// Mock Azure serverless URL
128+
DatabaseMetaData mockDatabaseMetaData = Mockito.mock(DatabaseMetaData.class);
129+
Mockito.when(connection.getMetaData()).thenReturn(mockDatabaseMetaData);
130+
Mockito.when(mockDatabaseMetaData.getURL()).thenReturn("datalakegentwo://jdbc:sqlserver://myworkspace-ondemand.sql.azuresynapse.net:1433;database=mydatabase;");
131+
132+
// Mock the prepared statement and result set
133+
PreparedStatement mockPreparedStatement = Mockito.mock(PreparedStatement.class);
134+
java.sql.ResultSet mockResultSet = Mockito.mock(java.sql.ResultSet.class);
135+
Mockito.when(connection.prepareStatement(Mockito.anyString())).thenReturn(mockPreparedStatement);
136+
Mockito.when(mockPreparedStatement.executeQuery()).thenReturn(mockResultSet);
137+
Mockito.when(mockResultSet.next()).thenReturn(false); // No rows
138+
139+
// Mock the query status checker
140+
com.amazonaws.athena.connector.lambda.QueryStatusChecker mockQueryStatusChecker =
141+
Mockito.mock(com.amazonaws.athena.connector.lambda.QueryStatusChecker.class);
142+
Mockito.when(mockQueryStatusChecker.isQueryRunning()).thenReturn(true);
143+
144+
// Mock the block spiller
145+
com.amazonaws.athena.connector.lambda.data.BlockSpiller mockBlockSpiller =
146+
Mockito.mock(com.amazonaws.athena.connector.lambda.data.BlockSpiller.class);
147+
148+
// Create the read records request
149+
com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest mockReadRecordsRequest =
150+
Mockito.mock(com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest.class);
151+
Mockito.when(mockReadRecordsRequest.getQueryId()).thenReturn("testQueryId");
152+
Mockito.when(mockReadRecordsRequest.getCatalogName()).thenReturn("testCatalog");
153+
Mockito.when(mockReadRecordsRequest.getTableName()).thenReturn(new TableName("testSchema", "testTable"));
154+
Mockito.when(mockReadRecordsRequest.getSchema()).thenReturn(SchemaBuilder.newBuilder().build());
155+
Mockito.when(mockReadRecordsRequest.getConstraints()).thenReturn(Mockito.mock(Constraints.class));
156+
Mockito.when(mockReadRecordsRequest.getSplit()).thenReturn(Mockito.mock(Split.class));
157+
158+
// Execute the method
159+
dataLakeGen2RecordHandler.readWithConstraint(mockBlockSpiller, mockReadRecordsRequest, mockQueryStatusChecker);
160+
161+
// Verify that setAutoCommit(false) was NOT called for Azure serverless environment
162+
Mockito.verify(connection, Mockito.never()).setAutoCommit(false);
163+
164+
// Verify that commit() was NOT called for Azure serverless environment
165+
Mockito.verify(connection, Mockito.never()).commit();
166+
}
167+
168+
@Test
169+
public void testReadWithConstraintWithStandardEnvironment()
170+
throws Exception
171+
{
172+
// Mock standard SQL Server URL (non-serverless)
173+
DatabaseMetaData mockDatabaseMetaData = Mockito.mock(DatabaseMetaData.class);
174+
Mockito.when(connection.getMetaData()).thenReturn(mockDatabaseMetaData);
175+
Mockito.when(mockDatabaseMetaData.getURL()).thenReturn("datalakegentwo://jdbc:sqlserver://myserver.database.windows.net:1433;database=mydatabase;");
176+
177+
// Mock the prepared statement and result set
178+
PreparedStatement mockPreparedStatement = Mockito.mock(PreparedStatement.class);
179+
java.sql.ResultSet mockResultSet = Mockito.mock(java.sql.ResultSet.class);
180+
Mockito.when(connection.prepareStatement(Mockito.anyString())).thenReturn(mockPreparedStatement);
181+
Mockito.when(mockPreparedStatement.executeQuery()).thenReturn(mockResultSet);
182+
Mockito.when(mockResultSet.next()).thenReturn(false); // No rows
183+
184+
// Mock the query status checker
185+
com.amazonaws.athena.connector.lambda.QueryStatusChecker mockQueryStatusChecker =
186+
Mockito.mock(com.amazonaws.athena.connector.lambda.QueryStatusChecker.class);
187+
Mockito.when(mockQueryStatusChecker.isQueryRunning()).thenReturn(true);
188+
189+
// Mock the block spiller
190+
com.amazonaws.athena.connector.lambda.data.BlockSpiller mockBlockSpiller =
191+
Mockito.mock(com.amazonaws.athena.connector.lambda.data.BlockSpiller.class);
192+
193+
// Create the read records request
194+
com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest mockReadRecordsRequest =
195+
Mockito.mock(com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest.class);
196+
Mockito.when(mockReadRecordsRequest.getQueryId()).thenReturn("testQueryId");
197+
Mockito.when(mockReadRecordsRequest.getCatalogName()).thenReturn("testCatalog");
198+
Mockito.when(mockReadRecordsRequest.getTableName()).thenReturn(new TableName("testSchema", "testTable"));
199+
Mockito.when(mockReadRecordsRequest.getSchema()).thenReturn(SchemaBuilder.newBuilder().build());
200+
Mockito.when(mockReadRecordsRequest.getConstraints()).thenReturn(Mockito.mock(Constraints.class));
201+
Mockito.when(mockReadRecordsRequest.getSplit()).thenReturn(Mockito.mock(Split.class));
202+
203+
// Execute the method
204+
dataLakeGen2RecordHandler.readWithConstraint(mockBlockSpiller, mockReadRecordsRequest, mockQueryStatusChecker);
205+
206+
// Verify that setAutoCommit(false) WAS called for standard environment
207+
Mockito.verify(connection, Mockito.times(1)).setAutoCommit(false);
208+
209+
// Verify that commit() WAS called for standard environment
210+
Mockito.verify(connection, Mockito.times(1)).commit();
211+
}
122212
}

0 commit comments

Comments
 (0)