Skip to content

Commit 5c0ab48

Browse files
ritiktrianzritiktrianz
committed
added unit tests for athena-federation-sdk.
ild issue Add unit tests for security package Address Review Comments Added unit test class for S3BlockSpillReader
1 parent b04a070 commit 5c0ab48

22 files changed

+2961
-4
lines changed

‎athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/connection/EnvironmentPropertiesTest.java‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
/*-
2+
* #%L
3+
* Amazon Athena Query Federation SDK
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.connector.lambda.connection;
21+
22+
import org.junit.Before;
23+
import org.junit.Test;
24+
import software.amazon.awssdk.services.glue.model.AuthenticationConfiguration;
25+
import software.amazon.awssdk.services.glue.model.Connection;
26+
import software.amazon.awssdk.services.glue.model.ConnectionPropertyKey;
27+
28+
import java.util.HashMap;
29+
import java.util.Map;
30+
31+
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.DATABASE;
32+
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.DEFAULT_GLUE_CONNECTION;
33+
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.KMS_KEY_ID;
34+
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.SECRET_NAME;
35+
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.SPILL_KMS_KEY_ID;
36+
import static org.junit.jupiter.api.Assertions.assertEquals;
37+
import static org.mockito.Mockito.doReturn;
38+
import static org.mockito.Mockito.spy;
39+
40+
public class EnvironmentPropertiesTest {
41+
42+
// Given
43+
private String glueConnName = "my-glue-conn";
44+
private String secretArn = "arn:aws:secretsmanager:us-east-1:1234567890:secret:my-secret-abc123";
45+
private String expectedSecretName = "my-secret";
46+
private String testValue = "test";
47+
private String kmsKeyId = "kms-123";
48+
private String lambdaValue = "lambda-value";
49+
50+
@Before
51+
public void setUp()
52+
throws Exception {
53+
// Mock env
54+
Map<String, String> environmentVariables = new HashMap<>();
55+
environmentVariables.put(DEFAULT_GLUE_CONNECTION, glueConnName);
56+
environmentVariables.put("OVERRIDE_VAR", lambdaValue); // Simulate Lambda-provided env var
57+
58+
setEnvironmentVariables(environmentVariables);
59+
}
60+
61+
@Test
62+
public void testCreateEnvironment() {
63+
64+
// Create a partial mock so we can stub getGlueConnection
65+
EnvironmentProperties spyProps = spy(new EnvironmentProperties());
66+
67+
// Mock Glue connection
68+
AuthenticationConfiguration authConfig = AuthenticationConfiguration.builder()
69+
.secretArn(secretArn)
70+
.build();
71+
72+
Map<ConnectionPropertyKey, String> connectionProps = new HashMap<>();
73+
connectionProps.put(ConnectionPropertyKey.DATABASE, testValue);
74+
75+
Map<String, String> athenaProps = new HashMap<>();
76+
athenaProps.put(SPILL_KMS_KEY_ID, kmsKeyId);
77+
78+
Connection glueConnection = Connection.builder()
79+
.name(glueConnName)
80+
.connectionProperties(connectionProps)
81+
.authenticationConfiguration(authConfig)
82+
.athenaProperties(athenaProps)
83+
.build();
84+
85+
doReturn(glueConnection).when(spyProps).getGlueConnection(glueConnName);
86+
87+
Map<String, String> result = spyProps.createEnvironment();
88+
89+
assertEquals(glueConnName, result.get(DEFAULT_GLUE_CONNECTION));
90+
assertEquals(testValue, result.get(DATABASE));
91+
assertEquals(expectedSecretName, result.get(SECRET_NAME));
92+
assertEquals(kmsKeyId, result.get(KMS_KEY_ID));
93+
assertEquals(lambdaValue, result.get("OVERRIDE_VAR"));
94+
}
95+
96+
// Utility to override environment variables in tests
97+
private static void setEnvironmentVariables(Map<String, String> newEnvironmentVariables) {
98+
try {
99+
Map<String, String> env = System.getenv();
100+
Class<?> cl = env.getClass();
101+
java.lang.reflect.Field field = cl.getDeclaredField("m");
102+
field.setAccessible(true);
103+
Map<String, String> writableEnv = (Map<String, String>) field.get(env);
104+
writableEnv.clear();
105+
writableEnv.putAll(newEnvironmentVariables);
106+
} catch (Exception e) {
107+
throw new RuntimeException("Failed to set environment variables", e);
108+
}
109+
}
110+
}

‎athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/data/ArrowTypeComparatorTest.java‎

Lines changed: 245 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,245 @@
1+
/*-
2+
* #%L
3+
* Amazon Athena Query Federation SDK
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.connector.lambda.data;
21+
22+
import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
23+
import org.apache.arrow.vector.types.DateUnit;
24+
import org.apache.arrow.vector.types.FloatingPointPrecision;
25+
import org.apache.arrow.vector.types.IntervalUnit;
26+
import org.apache.arrow.vector.types.TimeUnit;
27+
import org.apache.arrow.vector.types.pojo.ArrowType;
28+
29+
import org.junit.Before;
30+
import org.junit.Test;
31+
32+
import java.math.BigDecimal;
33+
import java.time.LocalDateTime;
34+
import java.time.ZoneOffset;
35+
import java.util.HashMap;
36+
import java.util.Map;
37+
38+
import static org.junit.Assert.assertEquals;
39+
import static org.junit.Assert.assertThrows;
40+
41+
public class ArrowTypeComparatorTest {
42+
43+
private String stringABC = "abc";
44+
private String stringXYZ = "xyz";
45+
private BigDecimal decimal1 = new BigDecimal("123.45");
46+
private BigDecimal decimal2 = new BigDecimal("678.90");
47+
private byte[] binary1 = {0x01, 0x02};
48+
private byte[] binary2 = {0x01, 0x03};
49+
private LocalDateTime date1 = LocalDateTime.of(2023, 1, 1, 12, 0);
50+
private LocalDateTime date2 = LocalDateTime.of(2023, 1, 2, 12, 0);
51+
private Map<String, Object> struct1;
52+
private Map<String, Object> struct2;
53+
54+
private ArrowType intType = new ArrowType.Int(32, true); // Signed 32-bit integer
55+
private ArrowType tinyIntType = new ArrowType.Int(8, true); // Signed 8-bit integer
56+
private ArrowType smallIntType = new ArrowType.Int(16, true); // Signed 16-bit integer
57+
private ArrowType uint2Type = new ArrowType.Int(16, false); // Unsigned 16-bit integer
58+
private ArrowType bigIntType = new ArrowType.Int(64, true); // Signed 64-bit integer
59+
private ArrowType float8Type = new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE); // Double precision
60+
private ArrowType float4Type = new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE); // Single precision
61+
private ArrowType varcharType = new ArrowType.Utf8(); // VARCHAR
62+
private ArrowType varbinaryType = new ArrowType.Binary(); // VARBINARY
63+
private ArrowType decimalType = new ArrowType.Decimal(10, 2, 128); // DECIMAL
64+
private ArrowType bitType = new ArrowType.Bool(); // BOOLEAN
65+
private ArrowType dateMilliType = new ArrowType.Timestamp(TimeUnit.MILLISECOND, ZoneOffset.UTC.getId()); // TIMESTAMP
66+
private ArrowType dateDayType = new ArrowType.Date(DateUnit.DAY); // DATE
67+
private ArrowType structType = new ArrowType.Struct(); // STRUCT type
68+
private ArrowType unsupportedType = new ArrowType.Interval(IntervalUnit.YEAR_MONTH); // Unsupported interval type
69+
70+
@Before
71+
public void setUp() {
72+
struct1 = new HashMap<>();
73+
struct1.put("key1", "value1");
74+
struct1.put("key2", 42);
75+
76+
struct2 = new HashMap<>();
77+
struct2.put("key1", "value1");
78+
struct2.put("key2", 43); // Different value for "key2"
79+
}
80+
81+
private void assertComparison(int expected, ArrowType arrowType, Object lhs, Object rhs) {
82+
assertEquals(expected, ArrowTypeComparator.compare(arrowType, lhs, rhs));
83+
}
84+
85+
private void assertStructComparison(int expected, Map<String, Object> lhs, Map<String, Object> rhs) {
86+
assertEquals(expected, ArrowTypeComparator.compare(structType, lhs, rhs));
87+
}
88+
89+
@Test
90+
public void testCompareIntegers() {
91+
assertComparison(-1, intType, 5, 10);
92+
assertComparison(1, intType, 10, 5);
93+
assertComparison(0, intType, 5, 5);
94+
}
95+
96+
@Test
97+
public void testCompareTinyInt() {
98+
assertComparison(-1, tinyIntType, (byte) 1, (byte) 2);
99+
assertComparison(1, tinyIntType, (byte) 2, (byte) 1);
100+
assertComparison(0, tinyIntType, (byte) 1, (byte) 1);
101+
}
102+
103+
@Test
104+
public void testCompareSmallInt() {
105+
assertComparison(-100, smallIntType, (short) 100, (short) 200);
106+
assertComparison(100, smallIntType, (short) 200, (short) 100);
107+
assertComparison(0, smallIntType, (short) 100, (short) 100);
108+
}
109+
110+
@Test
111+
public void testCompareUint2() {
112+
assertComparison(-1, uint2Type, 'A', 'B');
113+
assertComparison(1, uint2Type, 'B', 'A');
114+
assertComparison(0, uint2Type, 'A', 'A');
115+
}
116+
117+
@Test
118+
public void testCompareBigInt() {
119+
assertComparison(-1, bigIntType, 1000L, 2000L);
120+
assertComparison(1, bigIntType, 2000L, 1000L);
121+
assertComparison(0, bigIntType, 1000L, 1000L);
122+
}
123+
124+
@Test
125+
public void testCompareFloat8() {
126+
assertComparison(-1, float8Type, 123.45, 678.90);
127+
assertComparison(1, float8Type, 678.90, 123.45);
128+
assertComparison(0, float8Type, 123.45, 123.45);
129+
}
130+
131+
@Test
132+
public void testCompareFloat4() {
133+
assertComparison(-1, float4Type, 123.45f, 678.90f);
134+
assertComparison(1, float4Type, 678.90f, 123.45f);
135+
assertComparison(0, float4Type, 123.45f, 123.45f);
136+
}
137+
138+
@Test
139+
public void testCompareVarchar() {
140+
assertComparison(-23, varcharType, stringABC, stringXYZ);
141+
assertComparison(23, varcharType, stringXYZ, stringABC);
142+
assertComparison(0, varcharType, stringABC, stringABC);
143+
}
144+
145+
@Test
146+
public void testCompareVarbinary() {
147+
assertComparison(-1, varbinaryType, binary1, binary2);
148+
assertComparison(1, varbinaryType, binary2, binary1);
149+
assertComparison(0, varbinaryType, binary1, binary1);
150+
}
151+
152+
@Test
153+
public void testCompareDecimal() {
154+
assertComparison(-1, decimalType, decimal1, decimal2);
155+
assertComparison(1, decimalType, decimal2, decimal1);
156+
assertComparison(0, decimalType, decimal1, decimal1);
157+
}
158+
159+
@Test
160+
public void testCompareBit() {
161+
assertComparison(1, bitType, true, false);
162+
assertComparison(-1, bitType, false, true);
163+
assertComparison(0, bitType, true, true);
164+
}
165+
166+
@Test
167+
public void testCompareDateMilli() {
168+
assertComparison(-1, dateMilliType, date1, date2);
169+
assertComparison(1, dateMilliType, date2, date1);
170+
assertComparison(0, dateMilliType, date1, date1);
171+
}
172+
173+
@Test
174+
public void testCompareDateDay() {
175+
assertComparison(-1, dateDayType, 1, 2);
176+
assertComparison(1, dateDayType, 2, 1);
177+
assertComparison(0, dateDayType, 1, 1);
178+
}
179+
180+
@Test
181+
public void testCompareStructEqual() {
182+
struct2.put("key2", 42);
183+
184+
assertStructComparison(0, struct1, struct2);
185+
}
186+
187+
@Test
188+
public void testCompareStructNotEqual() {
189+
// Act & Assert: Expect inequality
190+
int result = ArrowTypeComparator.compare(structType, struct1, struct2);
191+
if (struct1.hashCode() < struct2.hashCode()) {
192+
assertEquals(-1, result);
193+
} else {
194+
assertEquals(1, result);
195+
}
196+
}
197+
198+
@Test
199+
public void testCompareStructWithNull() {
200+
// Arrange: Set one struct to null
201+
struct2 = null;
202+
203+
// Act & Assert: Expect struct1 (non-null) to be greater than struct2 (null)
204+
assertStructComparison(-1, struct1, struct2);
205+
206+
// Act & Assert: Reverse comparison
207+
assertStructComparison(1, struct2, struct1);
208+
}
209+
210+
@Test
211+
public void testCompareStructBothNull() {
212+
// Arrange: Set both structs to null
213+
struct1 = null;
214+
struct2 = null;
215+
216+
// Act & Assert: Expect equality
217+
assertStructComparison(0, struct1, struct2);
218+
}
219+
220+
@Test
221+
public void testCompareUnsupportedStructType() {
222+
// Act & Assert: Expect AthenaConnectorException
223+
AthenaConnectorException exception = assertThrows(
224+
AthenaConnectorException.class,
225+
() -> ArrowTypeComparator.compare(unsupportedType, struct1, struct2)
226+
);
227+
assertEquals("Unknown type INTERVALYEAR object: class java.util.HashMap", exception.getMessage());
228+
}
229+
230+
@Test
231+
public void testUnsupportedType() {
232+
AthenaConnectorException exception = assertThrows(
233+
AthenaConnectorException.class,
234+
() -> ArrowTypeComparator.compare(unsupportedType, "lhs", "rhs")
235+
);
236+
assertEquals("Unknown type INTERVALYEAR object: class java.lang.String", exception.getMessage());
237+
}
238+
239+
@Test
240+
public void testCompareNulls() {
241+
assertComparison(0, intType, null, null);
242+
assertComparison(1, intType, null, 10);
243+
assertComparison(-1, intType, 10, null);
244+
}
245+
}

0 commit comments

Comments
 (0)