Skip to content

Commit 86081d8

Browse files
Trianz-AkshayVenkatasivareddyTRritiktrianzburhan94
authored
Added unit tests for JDBC module (#2732)
Co-authored-by: VenkatasivaReddy Avula <[email protected]> Co-authored-by: ritik <[email protected]> Co-authored-by: burhan94 <[email protected]>
1 parent b668b21 commit 86081d8

11 files changed

+2646
-138
lines changed

‎athena-jdbc/src/test/java/com/amazonaws/athena/connectors/jdbc/JdbcEnvironmentPropertiesTest.java‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
/*-
2+
* #%L
3+
* athena-jdbc
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.jdbc;
21+
22+
import org.junit.Before;
23+
import org.junit.Test;
24+
import java.util.HashMap;
25+
import java.util.Map;
26+
27+
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.*;
28+
import static org.junit.Assert.assertEquals;
29+
30+
public class JdbcEnvironmentPropertiesTest
31+
{
32+
private static final String TEST_HOST = "test.host.com";
33+
private static final String TEST_PORT = "1234";
34+
private static final String TEST_DATABASE = "testdb";
35+
private static final String TEST_SECRET = "testSecret";
36+
private static final String TEST_ENCRYPT_PARAM = "encrypt=false";
37+
private static final String CONNECTION_STRING_PREFIX = "databaseName://jdbc:databaseName://";
38+
private static final String BASE_CONNECTION_STRING = CONNECTION_STRING_PREFIX + TEST_HOST + ":" + TEST_PORT + "/" + TEST_DATABASE;
39+
40+
private Map<String, String> connectionProperties;
41+
private JdbcEnvironmentProperties jdbcEnvironmentProperties;
42+
43+
@Before
44+
public void setup()
45+
{
46+
connectionProperties = new HashMap<>();
47+
connectionProperties.put(HOST, TEST_HOST);
48+
connectionProperties.put(PORT, TEST_PORT);
49+
connectionProperties.put(DATABASE, TEST_DATABASE);
50+
jdbcEnvironmentProperties = new JdbcEnvironmentProperties() {
51+
@Override
52+
protected String getConnectionStringPrefix(Map<String, String> connectionProperties) {
53+
return CONNECTION_STRING_PREFIX;
54+
}
55+
};
56+
}
57+
58+
@Test
59+
public void testConnectionPropertiesWithNoParams() {
60+
Map<String, String> result = jdbcEnvironmentProperties.connectionPropertiesToEnvironment(connectionProperties);
61+
String expected = BASE_CONNECTION_STRING + "?";
62+
63+
assertEquals(expected, result.get(DEFAULT));
64+
}
65+
66+
@Test
67+
public void testConnectionPropertiesWithOnlySecret() {
68+
connectionProperties.put(SECRET_NAME, TEST_SECRET);
69+
Map<String, String> result = jdbcEnvironmentProperties.connectionPropertiesToEnvironment(connectionProperties);
70+
String expected = BASE_CONNECTION_STRING + "?${" + TEST_SECRET + "}";
71+
72+
assertEquals(expected, result.get(DEFAULT));
73+
}
74+
75+
@Test
76+
public void testConnectionPropertiesWithOnlyJdbcParams() {
77+
final String sslParam = "ssl=true";
78+
connectionProperties.put(JDBC_PARAMS, sslParam);
79+
Map<String, String> result = jdbcEnvironmentProperties.connectionPropertiesToEnvironment(connectionProperties);
80+
String expected = BASE_CONNECTION_STRING + "?" + sslParam;
81+
82+
assertEquals(expected, result.get(DEFAULT));
83+
}
84+
85+
@Test
86+
public void testConnectionPropertiesWithDatabaseAndParamsAndSecret() {
87+
connectionProperties.put(JDBC_PARAMS, TEST_ENCRYPT_PARAM);
88+
connectionProperties.put(SECRET_NAME, TEST_SECRET);
89+
Map<String, String> result = jdbcEnvironmentProperties.connectionPropertiesToEnvironment(connectionProperties);
90+
String expected = BASE_CONNECTION_STRING + "?" + TEST_ENCRYPT_PARAM + "&${" + TEST_SECRET + "}";
91+
92+
assertEquals(expected, result.get(DEFAULT));
93+
}
94+
95+
@Test
96+
public void testSpecialCharactersInConnectionParams() {
97+
String specialHost = "test.host-with_special.chars.com";
98+
String specialDb = "test/db#1";
99+
String specialParams = "param1=value1&param2=value2#hash;param3=value3";
100+
101+
connectionProperties.put(HOST, specialHost);
102+
connectionProperties.put(DATABASE, specialDb);
103+
connectionProperties.put(JDBC_PARAMS, specialParams);
104+
105+
Map<String, String> result = jdbcEnvironmentProperties.connectionPropertiesToEnvironment(connectionProperties);
106+
String expected = CONNECTION_STRING_PREFIX + specialHost + ":" + TEST_PORT + "/" + specialDb + "?" + specialParams;
107+
assertEquals(expected, result.get(DEFAULT));
108+
}
109+
110+
@Test
111+
public void testMultipleJdbcParameters() {
112+
String params = "ssl=true&connectTimeout=10000&socketTimeout=5000&autoReconnect=true&useUnicode=true&characterEncoding=UTF-8";
113+
connectionProperties.put(JDBC_PARAMS, params);
114+
115+
Map<String, String> result = jdbcEnvironmentProperties.connectionPropertiesToEnvironment(connectionProperties);
116+
String expected = BASE_CONNECTION_STRING + "?" + params;
117+
assertEquals(expected, result.get(DEFAULT));
118+
}
119+
120+
@Test
121+
public void testComplexDatabasePath() {
122+
String complexDb = "main/schema1/table2";
123+
connectionProperties.put(DATABASE, complexDb);
124+
125+
Map<String, String> result = jdbcEnvironmentProperties.connectionPropertiesToEnvironment(connectionProperties);
126+
String expected = CONNECTION_STRING_PREFIX + TEST_HOST + ":" + TEST_PORT + "/" + complexDb + "?";
127+
assertEquals(expected, result.get(DEFAULT));
128+
}
129+
130+
@Test
131+
public void testSecretWithSpecialCharacters() {
132+
String specialSecret = "secret/name#1";
133+
connectionProperties.put(SECRET_NAME, specialSecret);
134+
135+
Map<String, String> result = jdbcEnvironmentProperties.connectionPropertiesToEnvironment(connectionProperties);
136+
String expected = BASE_CONNECTION_STRING + "?${" + specialSecret + "}";
137+
assertEquals(expected, result.get(DEFAULT));
138+
}
139+
}

‎athena-jdbc/src/test/java/com/amazonaws/athena/connectors/jdbc/MultiplexingJdbcMetadataHandlerTest.java‎

Lines changed: 64 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@
2424
import com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl;
2525
import com.amazonaws.athena.connector.lambda.data.BlockWriter;
2626
import com.amazonaws.athena.connector.lambda.domain.TableName;
27+
import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
28+
import com.amazonaws.athena.connector.lambda.metadata.GetDataSourceCapabilitiesRequest;
2729
import com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest;
2830
import com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest;
2931
import com.amazonaws.athena.connector.lambda.metadata.GetTableRequest;
@@ -39,8 +41,11 @@
3941
import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient;
4042

4143
import java.util.Collections;
44+
import java.util.HashMap;
4245
import java.util.Map;
4346

47+
import static org.junit.jupiter.api.Assertions.assertThrows;
48+
import static org.junit.jupiter.api.Assertions.assertTrue;
4449
import static org.mockito.ArgumentMatchers.nullable;
4550

4651
public class MultiplexingJdbcMetadataHandlerTest
@@ -53,21 +58,27 @@ public class MultiplexingJdbcMetadataHandlerTest
5358
private AthenaClient athena;
5459
private QueryStatusChecker queryStatusChecker;
5560
private JdbcConnectionFactory jdbcConnectionFactory;
61+
private DatabaseConnectionConfig databaseConnectionConfig;
62+
63+
private static final String FAKE_DATABASE = "fakedatabase";
64+
private static final String TEST_CATALOG = "testCatalog";
65+
private static final String TEST_SECRET = "testSecret";
66+
private static final String UNSUPPORTED_CATALOG = "unsupportedCatalog";
67+
private static final String CONNECTION_STRING = FAKE_DATABASE + "://jdbc:" + FAKE_DATABASE + "://hostname/${" + TEST_SECRET + "}";
68+
private static final int MAX_CATALOGS = 100;
69+
private static final int TOO_MANY_CATALOGS = 101;
5670

5771
@Before
5872
public void setup()
5973
{
60-
//this.allocator = Mockito.mock(BlockAllocator.class);
6174
this.allocator = new BlockAllocatorImpl();
62-
//Mockito.when(this.allocator.createBlock(nullable(Schema.class))).thenReturn(Mockito.mock(Block.class));
6375
this.fakeDatabaseHandler = Mockito.mock(JdbcMetadataHandler.class);
64-
this.metadataHandlerMap = Collections.singletonMap("fakedatabase", this.fakeDatabaseHandler);
76+
this.metadataHandlerMap = Collections.singletonMap(FAKE_DATABASE, this.fakeDatabaseHandler);
6577
this.secretsManager = Mockito.mock(SecretsManagerClient.class);
6678
this.athena = Mockito.mock(AthenaClient.class);
6779
this.queryStatusChecker = Mockito.mock(QueryStatusChecker.class);
6880
this.jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class);
69-
DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", "fakedatabase",
70-
"fakedatabase://jdbc:fakedatabase://hostname/${testSecret}", "testSecret");
81+
databaseConnectionConfig = new DatabaseConnectionConfig(TEST_CATALOG, FAKE_DATABASE, CONNECTION_STRING, TEST_SECRET);
7182
this.jdbcMetadataHandler = new MultiplexingJdbcMetadataHandler(this.secretsManager, this.athena, this.jdbcConnectionFactory, this.metadataHandlerMap, databaseConnectionConfig, com.google.common.collect.ImmutableMap.of());
7283
}
7384

@@ -76,7 +87,7 @@ public void doListSchemaNames()
7687
throws Exception
7788
{
7889
ListSchemasRequest listSchemasRequest = Mockito.mock(ListSchemasRequest.class);
79-
Mockito.when(listSchemasRequest.getCatalogName()).thenReturn("fakedatabase");
90+
Mockito.when(listSchemasRequest.getCatalogName()).thenReturn(FAKE_DATABASE);
8091
this.jdbcMetadataHandler.doListSchemaNames(this.allocator, listSchemasRequest);
8192
Mockito.verify(this.fakeDatabaseHandler, Mockito.times(1)).doListSchemaNames(Mockito.eq(this.allocator), Mockito.eq(listSchemasRequest));
8293
}
@@ -86,7 +97,7 @@ public void doListTables()
8697
throws Exception
8798
{
8899
ListTablesRequest listTablesRequest = Mockito.mock(ListTablesRequest.class);
89-
Mockito.when(listTablesRequest.getCatalogName()).thenReturn("fakedatabase");
100+
Mockito.when(listTablesRequest.getCatalogName()).thenReturn(FAKE_DATABASE);
90101
this.jdbcMetadataHandler.doListTables(this.allocator, listTablesRequest);
91102
Mockito.verify(this.fakeDatabaseHandler, Mockito.times(1)).doListTables(Mockito.eq(this.allocator), Mockito.eq(listTablesRequest));
92103
}
@@ -96,7 +107,7 @@ public void doGetTable()
96107
throws Exception
97108
{
98109
GetTableRequest getTableRequest = Mockito.mock(GetTableRequest.class);
99-
Mockito.when(getTableRequest.getCatalogName()).thenReturn("fakedatabase");
110+
Mockito.when(getTableRequest.getCatalogName()).thenReturn(FAKE_DATABASE);
100111
this.jdbcMetadataHandler.doGetTable(this.allocator, getTableRequest);
101112
Mockito.verify(this.fakeDatabaseHandler, Mockito.times(1)).doGetTable(Mockito.eq(this.allocator), Mockito.eq(getTableRequest));
102113
}
@@ -107,22 +118,22 @@ public void doGetTableLayout()
107118
{
108119
GetTableLayoutRequest getTableLayoutRequest = Mockito.mock(GetTableLayoutRequest.class);
109120
Mockito.when(getTableLayoutRequest.getTableName()).thenReturn(new TableName("testSchema", "testTable"));
110-
Mockito.when(getTableLayoutRequest.getCatalogName()).thenReturn("fakedatabase");
121+
Mockito.when(getTableLayoutRequest.getCatalogName()).thenReturn(FAKE_DATABASE);
111122
this.jdbcMetadataHandler.doGetTableLayout(this.allocator, getTableLayoutRequest);
112123
Mockito.verify(this.fakeDatabaseHandler, Mockito.times(1)).doGetTableLayout(Mockito.eq(this.allocator), Mockito.eq(getTableLayoutRequest));
113124
}
114125

115126
@Test
116127
public void getPartitionSchema()
117128
{
118-
this.jdbcMetadataHandler.getPartitionSchema("fakedatabase");
119-
Mockito.verify(this.fakeDatabaseHandler, Mockito.times(1)).getPartitionSchema(Mockito.eq("fakedatabase"));
129+
this.jdbcMetadataHandler.getPartitionSchema(FAKE_DATABASE);
130+
Mockito.verify(this.fakeDatabaseHandler, Mockito.times(1)).getPartitionSchema(Mockito.eq(FAKE_DATABASE));
120131
}
121132

122133
@Test(expected = RuntimeException.class)
123134
public void getPartitionSchemaForUnsupportedCatalog()
124135
{
125-
this.jdbcMetadataHandler.getPartitionSchema("unsupportedCatalog");
136+
this.jdbcMetadataHandler.getPartitionSchema(UNSUPPORTED_CATALOG);
126137
}
127138

128139

@@ -131,7 +142,7 @@ public void getPartitions()
131142
throws Exception
132143
{
133144
GetTableLayoutRequest getTableLayoutRequest = Mockito.mock(GetTableLayoutRequest.class);
134-
Mockito.when(getTableLayoutRequest.getCatalogName()).thenReturn("fakedatabase");
145+
Mockito.when(getTableLayoutRequest.getCatalogName()).thenReturn(FAKE_DATABASE);
135146
this.jdbcMetadataHandler.getPartitions(Mockito.mock(BlockWriter.class), getTableLayoutRequest, queryStatusChecker);
136147
Mockito.verify(this.fakeDatabaseHandler, Mockito.times(1)).getPartitions(nullable(BlockWriter.class), Mockito.eq(getTableLayoutRequest), Mockito.eq(queryStatusChecker));
137148
}
@@ -140,8 +151,47 @@ public void getPartitions()
140151
public void doGetSplits()
141152
{
142153
GetSplitsRequest getSplitsRequest = Mockito.mock(GetSplitsRequest.class);
143-
Mockito.when(getSplitsRequest.getCatalogName()).thenReturn("fakedatabase");
154+
Mockito.when(getSplitsRequest.getCatalogName()).thenReturn(FAKE_DATABASE);
144155
this.jdbcMetadataHandler.doGetSplits(this.allocator, getSplitsRequest);
145156
Mockito.verify(this.fakeDatabaseHandler, Mockito.times(1)).doGetSplits(Mockito.eq(this.allocator), Mockito.eq(getSplitsRequest));
146157
}
158+
159+
@Test
160+
public void testConstructor_withTooManyHandlers_shouldThrowException() {
161+
metadataHandlerMap = new HashMap<>();
162+
for (int i = 0; i < TOO_MANY_CATALOGS; i++) {
163+
metadataHandlerMap.put("catalog" + i, fakeDatabaseHandler);
164+
}
165+
166+
AthenaConnectorException exception = assertThrows(AthenaConnectorException.class, () ->
167+
new MultiplexingJdbcMetadataHandler(
168+
secretsManager,
169+
athena,
170+
jdbcConnectionFactory,
171+
metadataHandlerMap,
172+
databaseConnectionConfig,
173+
com.google.common.collect.ImmutableMap.of()
174+
)
175+
);
176+
assertTrue(exception.getMessage().contains("Max " + MAX_CATALOGS + " catalogs supported in multiplexer"));
177+
}
178+
179+
@Test
180+
public void testDoGetQueryPassthroughSchema()
181+
throws Exception
182+
{
183+
GetTableRequest getTableRequest = Mockito.mock(GetTableRequest.class);
184+
Mockito.when(getTableRequest.getCatalogName()).thenReturn(FAKE_DATABASE);
185+
this.jdbcMetadataHandler.doGetQueryPassthroughSchema(this.allocator, getTableRequest);
186+
Mockito.verify(this.fakeDatabaseHandler, Mockito.times(1)).doGetQueryPassthroughSchema(Mockito.eq(this.allocator), Mockito.eq(getTableRequest));
187+
}
188+
189+
@Test
190+
public void testDoGetDataSourceCapabilities()
191+
{
192+
GetDataSourceCapabilitiesRequest getDataSourceCapabilitiesRequest = Mockito.mock(GetDataSourceCapabilitiesRequest.class);
193+
Mockito.when(getDataSourceCapabilitiesRequest.getCatalogName()).thenReturn(FAKE_DATABASE);
194+
this.jdbcMetadataHandler.doGetDataSourceCapabilities(this.allocator, getDataSourceCapabilitiesRequest);
195+
Mockito.verify(this.fakeDatabaseHandler, Mockito.times(1)).doGetDataSourceCapabilities(Mockito.eq(this.allocator), Mockito.eq(getDataSourceCapabilitiesRequest));
196+
}
147197
}

0 commit comments

Comments
 (0)