Skip to content

Commit 4c600d4

Browse files
ritiktrianzJithendar12
ritiktrianz
authored and
Jithendar12
committed
Add unit tests for athena-oracle. (awslabs#2836)
Authored-by: Jithendar12 <[email protected]>
1 parent bc06988 commit 4c600d4

File tree

8 files changed

+937
-94
lines changed

8 files changed

+937
-94
lines changed

‎athena-oracle/src/main/java/com/amazonaws/athena/connectors/oracle/OracleJdbcConnectionFactory.java‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,13 +25,15 @@
2525
import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig;
2626
import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionInfo;
2727
import com.amazonaws.athena.connectors.jdbc.connection.GenericJdbcConnectionFactory;
28+
import com.google.common.annotations.VisibleForTesting;
2829
import org.apache.commons.lang3.Validate;
2930
import org.slf4j.Logger;
3031
import org.slf4j.LoggerFactory;
3132

3233
import java.sql.Connection;
3334
import java.sql.DriverManager;
3435
import java.sql.SQLException;
36+
import java.util.Map;
3537
import java.util.Properties;
3638
import java.util.regex.Matcher;
3739

@@ -71,7 +73,7 @@ public Connection getConnection(final CredentialsProvider credentialsProvider)
7173
// By default; Oracle RDS uses SSL_RSA_WITH_AES_256_CBC_SHA
7274
// Adding the following cipher suits to support others listed in Doc
7375
// https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/Appendix.Oracle.Options.SSL.html#Appendix.Oracle.Options.SSL.CipherSuites
74-
if (System.getenv().getOrDefault(IS_FIPS_ENABLED, "false").equalsIgnoreCase("true") || System.getenv().getOrDefault(IS_FIPS_ENABLED_LEGACY, "false").equalsIgnoreCase("true")) {
76+
if (getEnvMap().getOrDefault(IS_FIPS_ENABLED, "false").equalsIgnoreCase("true") || getEnvMap().getOrDefault(IS_FIPS_ENABLED_LEGACY, "false").equalsIgnoreCase("true")) {
7577
properties.put("oracle.net.ssl_cipher_suites", "(TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384, TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256, TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA384, TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA256, TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA, TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA)");
7678
}
7779
}
@@ -101,4 +103,13 @@ public Connection getConnection(final CredentialsProvider credentialsProvider)
101103
throw new RuntimeException(ex);
102104
}
103105
}
106+
107+
/**
108+
* Extracted method for environment variables to allow overriding in tests.
109+
*/
110+
@VisibleForTesting
111+
protected Map<String, String> getEnvMap()
112+
{
113+
return System.getenv();
114+
}
104115
}

‎athena-oracle/src/test/java/com/amazonaws/athena/connectors/oracle/OracleEnvironmentPropertiesTest.java‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
/*-
2+
* #%L
3+
* athena-oracle
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.oracle;
21+
22+
import org.junit.Before;
23+
import org.junit.Test;
24+
25+
import java.util.HashMap;
26+
import java.util.Map;
27+
28+
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.DATABASE;
29+
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.DEFAULT;
30+
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.ENFORCE_SSL;
31+
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.HOST;
32+
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.PORT;
33+
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.SECRET_NAME;
34+
import static org.junit.Assert.assertEquals;
35+
36+
public class OracleEnvironmentPropertiesTest
37+
{
38+
Map<String, String> connectionProperties;
39+
OracleEnvironmentProperties oracleEnvironmentProperties;
40+
41+
@Before
42+
public void setUp()
43+
{
44+
connectionProperties = new HashMap<>();
45+
connectionProperties.put(HOST, "test.oracle.com");
46+
connectionProperties.put(PORT, "1521");
47+
connectionProperties.put(DATABASE, "orcl");
48+
connectionProperties.put(SECRET_NAME, "oracle-secret");
49+
oracleEnvironmentProperties = new OracleEnvironmentProperties();
50+
}
51+
52+
@Test
53+
public void connectionPropertiesToEnvironment_withoutSSLEnabled_returnsConnectionStringWithoutSSL()
54+
{
55+
Map<String, String> oracleConnectionProperties = oracleEnvironmentProperties.connectionPropertiesToEnvironment(connectionProperties);
56+
String expectedConnectionString = "oracle://jdbc:oracle:thin:${oracle-secret}@//test.oracle.com:1521/orcl";
57+
assertEquals(expectedConnectionString, oracleConnectionProperties.get(DEFAULT));
58+
}
59+
60+
@Test
61+
public void connectionPropertiesToEnvironment_withSSLEnabled_returnsConnectionStringWithSSL()
62+
{
63+
connectionProperties.put(ENFORCE_SSL, "true");
64+
Map<String, String> oracleConnectionProperties = oracleEnvironmentProperties.connectionPropertiesToEnvironment(connectionProperties);
65+
String expectedConnectionString = "oracle://jdbc:oracle:thin:${oracle-secret}@tcps://test.oracle.com:1521/orcl";
66+
assertEquals(expectedConnectionString, oracleConnectionProperties.get(DEFAULT));
67+
}
68+
}

‎athena-oracle/src/test/java/com/amazonaws/athena/connectors/oracle/OracleJdbcConnectionFactoryTest.java‎

Lines changed: 123 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -19,33 +19,140 @@
1919
*/
2020
package com.amazonaws.athena.connectors.oracle;
2121

22-
import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig;
23-
import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionInfo;
24-
import com.amazonaws.athena.connector.credentials.DefaultCredentials;
2522
import com.amazonaws.athena.connector.credentials.CredentialsProvider;
23+
import com.amazonaws.athena.connector.credentials.DefaultCredentials;
2624
import com.amazonaws.athena.connector.credentials.StaticCredentialsProvider;
27-
import org.junit.Assert;
25+
import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig;
26+
import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionInfo;
2827
import org.junit.Test;
2928

3029
import java.sql.Connection;
3130
import java.sql.Driver;
3231
import java.sql.DriverManager;
3332
import java.sql.SQLException;
33+
import java.util.HashMap;
34+
import java.util.Map;
35+
import java.util.Properties;
3436

35-
public class OracleJdbcConnectionFactoryTest {
37+
import static org.junit.Assert.assertEquals;
38+
import static org.mockito.ArgumentMatchers.any;
39+
import static org.mockito.ArgumentMatchers.anyString;
40+
import static org.mockito.Mockito.mock;
41+
import static org.mockito.Mockito.when;
3642

37-
@Test(expected = RuntimeException.class)
38-
public void getConnectionTest() throws SQLException {
39-
DefaultCredentials expectedCredential = new DefaultCredentials("test", "test");
43+
public class OracleJdbcConnectionFactoryTest
44+
{
45+
private static final String CONNECTION_STRING = "oracle://jdbc:oracle:thin:username/password@//test.oracle.com:1521/orcl";
46+
private static final String TEST_CATALOG = "testCatalog";
47+
private static final String TEST_SECRET = "test";
48+
private static final String CATALOG = "catalog";
49+
private static final String USERNAME = "user";
50+
private static final String PASSWORD = "password";
51+
52+
@Test
53+
public void getConnection_withSslConnection_setsSslProperties() throws Exception
54+
{
55+
Driver mockDriver = mock(Driver.class);
56+
57+
DefaultCredentials expectedCredential = new DefaultCredentials(USERNAME, PASSWORD);
4058
CredentialsProvider credentialsProvider = new StaticCredentialsProvider(expectedCredential);
41-
DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", OracleConstants.ORACLE_NAME,
42-
"oracle://jdbc:oracle:thin:username/password@//127.0.0.1:1521/orcl", "test");
43-
DatabaseConnectionInfo DatabaseConnectionInfo = new DatabaseConnectionInfo(OracleConstants.ORACLE_DRIVER_CLASS,OracleConstants.ORACLE_DEFAULT_PORT);
44-
Connection connection = new OracleJdbcConnectionFactory(databaseConnectionConfig, DatabaseConnectionInfo).getConnection(credentialsProvider);
45-
String originalURL = connection.getMetaData().getURL();
46-
Driver drv = DriverManager.getDriver(originalURL);
47-
String driverClass = drv.getClass().getName();
48-
Assert.assertEquals("oracle.jdbc.OracleDriver", driverClass);
59+
60+
DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig(
61+
TEST_CATALOG,
62+
OracleConstants.ORACLE_NAME,
63+
"oracle://jdbc:oracle:thin:username/password@tcps://test.oracle.com:1521/orcl",
64+
TEST_SECRET
65+
);
66+
67+
DatabaseConnectionInfo databaseConnectionInfo = new DatabaseConnectionInfo(
68+
OracleConstants.ORACLE_DRIVER_CLASS,
69+
OracleConstants.ORACLE_DEFAULT_PORT
70+
);
71+
72+
Properties[] capturedProps = new Properties[1];
73+
74+
when(mockDriver.acceptsURL(anyString())).thenReturn(true);
75+
when(mockDriver.connect(anyString(), any(Properties.class)))
76+
.thenAnswer(invocation -> {
77+
capturedProps[0] = invocation.getArgument(1);
78+
return mock(Connection.class);
79+
});
80+
81+
try {
82+
DriverManager.registerDriver(mockDriver);
83+
84+
// Create a test connection factory with custom environment
85+
OracleJdbcConnectionFactory connectionFactory = new OracleJdbcConnectionFactory(databaseConnectionConfig, databaseConnectionInfo) {
86+
@Override
87+
protected Map<String, String> getEnvMap()
88+
{
89+
Map<String, String> env = new HashMap<>();
90+
env.put(IS_FIPS_ENABLED, "true");
91+
return env;
92+
}
93+
};
94+
95+
connectionFactory.getConnection(credentialsProvider);
96+
97+
Properties sslProps = capturedProps[0];
98+
assertEquals("JKS", sslProps.getProperty("javax.net.ssl.trustStoreType"));
99+
assertEquals("changeit", sslProps.getProperty("javax.net.ssl.trustStorePassword"));
100+
assertEquals("true", sslProps.getProperty("oracle.net.ssl_server_dn_match"));
101+
assertEquals(
102+
"(TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384, TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256, TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA384, TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA256, TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA, TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA)",
103+
sslProps.getProperty("oracle.net.ssl_cipher_suites")
104+
);
105+
}
106+
finally {
107+
DriverManager.deregisterDriver(mockDriver);
108+
}
109+
}
110+
111+
@Test(expected = RuntimeException.class)
112+
public void getConnection_withNullCredentialsProvider_throwsRuntimeException()
113+
{
114+
DatabaseConnectionConfig config = new DatabaseConnectionConfig(
115+
TEST_CATALOG, OracleConstants.ORACLE_NAME,
116+
CONNECTION_STRING, TEST_SECRET);
117+
DatabaseConnectionInfo info = new DatabaseConnectionInfo(OracleConstants.ORACLE_DRIVER_CLASS, OracleConstants.ORACLE_DEFAULT_PORT);
118+
119+
new OracleJdbcConnectionFactory(config, info).getConnection(null);
120+
}
121+
122+
@Test(expected = RuntimeException.class)
123+
public void getConnection_whenDriverFailsToConnect_throwsRuntimeException() throws Exception
124+
{
125+
DefaultCredentials credentials = new DefaultCredentials(USERNAME, PASSWORD);
126+
CredentialsProvider credentialsProvider = new StaticCredentialsProvider(credentials);
127+
128+
DatabaseConnectionConfig config = new DatabaseConnectionConfig(CATALOG, OracleConstants.ORACLE_NAME, CONNECTION_STRING, TEST_SECRET);
129+
DatabaseConnectionInfo info = new DatabaseConnectionInfo(OracleConstants.ORACLE_DRIVER_CLASS, OracleConstants.ORACLE_DEFAULT_PORT);
130+
131+
Driver mockDriver = mock(Driver.class);
132+
when(mockDriver.acceptsURL(anyString())).thenReturn(true);
133+
when(mockDriver.connect(anyString(), any(Properties.class)))
134+
.thenThrow(new SQLException("DB down", "08001", 1234));
135+
136+
try {
137+
DriverManager.registerDriver(mockDriver);
138+
new OracleJdbcConnectionFactory(config, info).getConnection(credentialsProvider);
139+
}
140+
finally {
141+
DriverManager.deregisterDriver(mockDriver);
142+
}
49143
}
50144

145+
@Test(expected = RuntimeException.class)
146+
public void getConnection_withInvalidDriverClass_throwsRuntimeException()
147+
{
148+
DefaultCredentials credentials = new DefaultCredentials(USERNAME, PASSWORD);
149+
CredentialsProvider credentialsProvider = new StaticCredentialsProvider(credentials);
150+
151+
DatabaseConnectionConfig config = new DatabaseConnectionConfig(CATALOG, OracleConstants.ORACLE_NAME, CONNECTION_STRING, TEST_SECRET);
152+
153+
// Provide an invalid/non-existent driver class name to trigger ClassNotFoundException
154+
DatabaseConnectionInfo info = new DatabaseConnectionInfo("non.existent.DriverClass", OracleConstants.ORACLE_DEFAULT_PORT);
155+
156+
new OracleJdbcConnectionFactory(config, info).getConnection(credentialsProvider);
157+
}
51158
}

0 commit comments

Comments
 (0)