Skip to content

Commit 41ff391

Browse files
VenkatasivareddyTRritiktrianzJithendar12
authored
added unit tests in athena-federation-sdk. (#2715)
Co-authored-by: ritik <[email protected]> Co-authored-by: Jithendar12 <[email protected]>
1 parent 3015ac8 commit 41ff391

29 files changed

+4492
-12
lines changed

‎athena-federation-sdk/pom.xml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,12 @@
5353
</exclusion>
5454
</exclusions>
5555
</dependency>
56+
<dependency>
57+
<groupId>com.github.stefanbirkner</groupId>
58+
<artifactId>system-lambda</artifactId>
59+
<version>1.2.1</version>
60+
<scope>test</scope>
61+
</dependency>
5662
<dependency>
5763
<groupId>software.amazon.awssdk</groupId>
5864
<artifactId>kms</artifactId>

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

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
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.Test;
23+
import software.amazon.awssdk.services.glue.model.AuthenticationConfiguration;
24+
import software.amazon.awssdk.services.glue.model.Connection;
25+
import software.amazon.awssdk.services.glue.model.ConnectionPropertyKey;
26+
27+
import java.util.HashMap;
28+
import java.util.Map;
29+
30+
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.DATABASE;
31+
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.DEFAULT_GLUE_CONNECTION;
32+
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.KMS_KEY_ID;
33+
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.SECRET_NAME;
34+
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.SPILL_KMS_KEY_ID;
35+
36+
import static org.junit.Assert.assertThrows;
37+
import static org.junit.jupiter.api.Assertions.assertEquals;
38+
import static com.github.stefanbirkner.systemlambda.SystemLambda.withEnvironmentVariable;
39+
import static org.mockito.Mockito.doReturn;
40+
import static org.mockito.Mockito.doThrow;
41+
import static org.mockito.Mockito.spy;
42+
43+
44+
public class EnvironmentPropertiesTest {
45+
46+
private static final String glueConnName = "my-glue-conn";
47+
private static final String secretArn = "arn:aws:secretsmanager:us-east-1:1234567890:secret:my-secret-abc123";
48+
private static final String expectedSecretName = "my-secret";
49+
private static final String testValue = "test";
50+
private static final String kmsKeyId = "kms-123";
51+
private static final String lambdaValue = "lambda-value";
52+
53+
54+
@Test
55+
public void testCreateEnvironmentWithSystemLambda() throws Exception {
56+
withEnvironmentVariable(DEFAULT_GLUE_CONNECTION, glueConnName)
57+
.and("OVERRIDE_VAR", lambdaValue)
58+
.execute(() -> {
59+
EnvironmentProperties spyProps = spy(new EnvironmentProperties());
60+
61+
AuthenticationConfiguration authConfig = AuthenticationConfiguration.builder()
62+
.secretArn(secretArn)
63+
.build();
64+
65+
Map<ConnectionPropertyKey, String> connectionProps = new HashMap<>();
66+
connectionProps.put(ConnectionPropertyKey.DATABASE, testValue);
67+
68+
Map<String, String> athenaProps = new HashMap<>();
69+
athenaProps.put(SPILL_KMS_KEY_ID, kmsKeyId);
70+
71+
Connection glueConnection = Connection.builder()
72+
.name(glueConnName)
73+
.connectionProperties(connectionProps)
74+
.authenticationConfiguration(authConfig)
75+
.athenaProperties(athenaProps)
76+
.build();
77+
78+
doReturn(glueConnection).when(spyProps).getGlueConnection(glueConnName);
79+
80+
Map<String, String> result = spyProps.createEnvironment();
81+
82+
assertEquals(glueConnName, result.get(DEFAULT_GLUE_CONNECTION));
83+
assertEquals(testValue, result.get(DATABASE));
84+
assertEquals(expectedSecretName, result.get(SECRET_NAME));
85+
assertEquals(kmsKeyId, result.get(KMS_KEY_ID));
86+
assertEquals(lambdaValue, result.get("OVERRIDE_VAR"));
87+
});
88+
}
89+
90+
@Test
91+
public void testCreateEnvironmentWithSystemLambda_GlueConnectionFails_ThrowsRuntimeException() throws Exception
92+
{
93+
withEnvironmentVariable(DEFAULT_GLUE_CONNECTION, glueConnName)
94+
.execute(() -> {
95+
EnvironmentProperties spyProps = spy(new EnvironmentProperties());
96+
97+
doThrow(new RuntimeException("Simulated failure"))
98+
.when(spyProps).getGlueConnection(glueConnName);
99+
100+
RuntimeException thrown = assertThrows(RuntimeException.class, spyProps::createEnvironment);
101+
102+
assertEquals("Simulated failure", thrown.getMessage());
103+
});
104+
}
105+
}

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

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

0 commit comments

Comments
 (0)