Skip to content

Commit 166606f

Browse files
committed
added unit test cases
1 parent 97ab1ad commit 166606f

File tree

4 files changed

+312
-37
lines changed

4 files changed

+312
-37
lines changed

athena-snowflake/src/main/java/com/amazonaws/athena/connectors/snowflake/SnowflakeQueryStringBuilder.java

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -248,7 +248,7 @@ protected String toPredicate(String columnName, String operator, Object value, A
248248
return columnName + " " + operator + " " + ((type.getTypeID().equals(Utf8) || type.getTypeID().equals(ArrowType.ArrowTypeID.Date)) ? singleQuote(getObjectForWhereClause(columnName, value, type).toString()) : getObjectForWhereClause(columnName, value, type));
249249
}
250250

251-
private static Object getObjectForWhereClause(String columnName, Object value, ArrowType arrowType)
251+
protected static Object getObjectForWhereClause(String columnName, Object value, ArrowType arrowType)
252252
{
253253
String val;
254254
StringBuilder tempVal;
@@ -279,12 +279,6 @@ else if (value instanceof Number) {
279279
DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
280280
return dateTime.format(formatter);
281281
}
282-
else if (val.contains("-")) {
283-
tempVal = new StringBuilder(val);
284-
tempVal = tempVal.length() == 19 ? tempVal.append(".0") : tempVal;
285-
val = String.format("%-26s", tempVal).replace(' ', '0').replace("T", " ");
286-
return val; // Returning as string formatted datetime
287-
}
288282
else {
289283
long days = Long.parseLong(val);
290284
long milliseconds = TimeUnit.DAYS.toMillis(days);
Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
package com.amazonaws.athena.connectors.snowflake;
2+
import static org.junit.jupiter.api.Assertions.*;
3+
import static org.mockito.Mockito.*;
4+
5+
import java.sql.Connection;
6+
import java.sql.PreparedStatement;
7+
import java.sql.ResultSet;
8+
import java.sql.SQLException;
9+
import java.util.HashMap;
10+
import java.util.Map;
11+
12+
import com.amazonaws.athena.connector.lambda.domain.TableName;
13+
import org.junit.jupiter.api.BeforeEach;
14+
import org.junit.jupiter.api.Test;
15+
import org.mockito.Mock;
16+
import org.mockito.Mockito;
17+
import org.mockito.MockitoAnnotations;
18+
19+
public class SnowflakeCaseInsensitiveResolverTest {
20+
21+
@Mock
22+
private Connection mockConnection;
23+
24+
@Mock
25+
private PreparedStatement mockPreparedStatement;
26+
27+
@Mock
28+
private ResultSet mockResultSet;
29+
30+
private Map<String, String> configOptions;
31+
32+
@BeforeEach
33+
void setUp() throws SQLException {
34+
MockitoAnnotations.openMocks(this);
35+
configOptions = new HashMap<>();
36+
37+
when(mockConnection.prepareStatement(anyString())).thenReturn(mockPreparedStatement);
38+
when(mockPreparedStatement.executeQuery()).thenReturn(mockResultSet);
39+
}
40+
//@Test
41+
void testCaseInsensitiveSearchMode() throws SQLException {
42+
configOptions.put("casing_mode", "CASE_INSENSITIVE_SEARCH");
43+
TableName inputTable = new TableName("TestSchema", "TestTable");
44+
45+
// Mock case-insensitive resolution
46+
SnowflakeCaseInsensitiveResolver spyResolver = Mockito.spy(SnowflakeCaseInsensitiveResolver.class);
47+
when(spyResolver.getSchemaNameCaseInsensitively(any(), any(), any())).thenReturn("testschema");
48+
when(spyResolver.getTableNameCaseInsensitively(any(), any(), any(), any())).thenReturn("testtable");
49+
50+
TableName result = SnowflakeCaseInsensitiveResolver.getAdjustedTableObjectNameBasedOnConfig(mockConnection, inputTable, configOptions);
51+
52+
assertEquals("testschema", result.getSchemaName());
53+
assertEquals("testtable", result.getTableName());
54+
}
55+
56+
@Test
57+
void testAnnotationMode() throws SQLException
58+
{
59+
configOptions.put("casing_mode", "ANNOTATION");
60+
TableName inputTable = new TableName("TestSchema", "TestTable@schemaCase=lower&tableCase=upper");
61+
62+
TableName result = SnowflakeCaseInsensitiveResolver.getAdjustedTableObjectNameBasedOnConfig(mockConnection, inputTable, configOptions);
63+
64+
assertEquals("testschema", result.getSchemaName());
65+
assertEquals("TESTTABLE", result.getTableName());
66+
}
67+
68+
// @Test
69+
void testNoneMode() throws SQLException
70+
{
71+
configOptions.put("casing_mode", "NONE");
72+
TableName inputTable = new TableName("TestSchema", "TestTable");
73+
74+
TableName result = SnowflakeCaseInsensitiveResolver.getAdjustedTableObjectNameBasedOnConfig(mockConnection, inputTable, configOptions);
75+
76+
assertEquals("TestSchema", result.getSchemaName());
77+
assertEquals("TestTable", result.getTableName());
78+
}
79+
80+
@Test
81+
void testInvalidCasingModeDefaultsToNone() {
82+
configOptions.put("casing_mode", "INVALID_MODE");
83+
TableName inputTable = new TableName("TestSchema", "TestTable");
84+
85+
assertThrows(Exception.class, () -> {
86+
SnowflakeCaseInsensitiveResolver.getAdjustedTableObjectNameBasedOnConfig(mockConnection, inputTable, configOptions);
87+
});
88+
}
89+
90+
@Test
91+
void testGetSchemaNameCaseInsensitively_Success() throws SQLException {
92+
when(mockResultSet.next()).thenReturn(true, false);
93+
when(mockResultSet.getString("SCHEMA_NAME")).thenReturn("TestSchema");
94+
95+
String result = SnowflakeCaseInsensitiveResolver.getSchemaNameCaseInsensitively(mockConnection, "testschema", configOptions);
96+
assertEquals("TestSchema", result);
97+
}
98+
99+
@Test
100+
void testGetSchemaNameCaseInsensitively_NoMatch() throws SQLException {
101+
when(mockResultSet.next()).thenReturn(false);
102+
103+
Exception exception = assertThrows(RuntimeException.class, () -> {
104+
SnowflakeCaseInsensitiveResolver.getSchemaNameCaseInsensitively(mockConnection, "unknownschema", configOptions);
105+
});
106+
107+
assertTrue(exception.getMessage().contains("Schema name case insensitive match failed"));
108+
}
109+
110+
@Test
111+
void testGetTableNameCaseInsensitively_Success() throws SQLException {
112+
when(mockResultSet.next()).thenReturn(true, false);
113+
when(mockResultSet.getString("TABLE_NAME")).thenReturn("TestTable");
114+
115+
String result = SnowflakeCaseInsensitiveResolver.getTableNameCaseInsensitively(mockConnection, "TestSchema", "testtable", configOptions);
116+
assertEquals("TestTable", result);
117+
}
118+
119+
@Test
120+
void testGetTableNameCaseInsensitively_NoMatch() throws SQLException {
121+
when(mockResultSet.next()).thenReturn(false);
122+
123+
Exception exception = assertThrows(RuntimeException.class, () -> {
124+
SnowflakeCaseInsensitiveResolver.getTableNameCaseInsensitively(mockConnection, "TestSchema", "unknownTable", configOptions);
125+
});
126+
127+
assertTrue(exception.getMessage().contains("Schema name case insensitive match failed"));
128+
}
129+
}

athena-snowflake/src/test/java/com/amazonaws/athena/connectors/snowflake/SnowflakeQueryStringBuilderTest.java

Lines changed: 115 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -18,41 +18,126 @@
1818
* #L%
1919
*/
2020
package com.amazonaws.athena.connectors.snowflake;
21-
import com.amazonaws.athena.connector.lambda.domain.Split;
22-
import org.junit.Assert;
23-
import org.junit.Test;
21+
22+
import com.amazonaws.athena.connector.lambda.domain.predicate.Constraints;
23+
import com.amazonaws.athena.connectors.jdbc.manager.TypeAndValue;
24+
import org.apache.arrow.vector.types.DateUnit;
25+
import org.apache.arrow.vector.types.TimeUnit;
26+
import org.apache.arrow.vector.types.pojo.ArrowType;
27+
import org.apache.arrow.vector.types.pojo.Field;
28+
import org.apache.arrow.vector.types.pojo.FieldType;
29+
import org.apache.arrow.vector.types.pojo.Schema;
30+
import org.junit.jupiter.api.BeforeEach;
31+
import org.junit.jupiter.api.Test;
2432
import org.mockito.Mock;
25-
import org.mockito.Mockito;
33+
import org.mockito.MockitoAnnotations;
2634

35+
import java.math.BigDecimal;
36+
import java.sql.Connection;
37+
import java.sql.SQLException;
2738
import java.util.ArrayList;
28-
import java.util.Collections;
2939
import java.util.List;
3040

31-
import static com.amazonaws.athena.connectors.snowflake.SnowflakeConstants.SNOWFLAKE_QUOTE_CHARACTER;
41+
import static org.junit.jupiter.api.Assertions.assertEquals;
42+
import static org.junit.jupiter.api.Assertions.assertThrows;
43+
import static org.junit.jupiter.api.Assertions.assertTrue;
44+
import static org.mockito.Mockito.mock;
45+
import static org.mockito.Mockito.when;
46+
47+
class SnowflakeQueryStringBuilderTest {
48+
49+
private SnowflakeQueryStringBuilder queryBuilder;
3250

33-
public class SnowflakeQueryStringBuilderTest
34-
{
3551
@Mock
36-
Split split;
37-
38-
//@Test
39-
// public void testQueryBuilderNew()
40-
// {
41-
// Split split = Mockito.mock(Split.class);
42-
// SnowflakeQueryStringBuilder builder = new SnowflakeQueryStringBuilder(SNOWFLAKE_QUOTE_CHARACTER, new SnowflakeFederationExpressionParser(SNOWFLAKE_QUOTE_CHARACTER));
43-
// Mockito.when(split.getProperties()).thenReturn(Collections.singletonMap("partition", "p0"));
44-
// Mockito.when(split.getProperty(Mockito.eq("partition"))).thenReturn("p1-p2-p3-p4-p5-p6-p7");
45-
// builder.getFromClauseWithSplit("default", "", "table", split);
46-
// builder.appendLimitOffset(split);
47-
// }
48-
49-
//@Test
50-
// public void testGetPartitionWhereClauses()
51-
// {
52-
// SnowflakeQueryStringBuilder builder = new SnowflakeQueryStringBuilder(SNOWFLAKE_QUOTE_CHARACTER, new SnowflakeFederationExpressionParser(SNOWFLAKE_QUOTE_CHARACTER));
53-
// List<String> fromClauseWithSplit = builder.getPartitionWhereClauses(split);
54-
// List<String> expected = new ArrayList<>();
55-
// Assert.assertEquals(expected, fromClauseWithSplit);
56-
// }
57-
}
52+
private Connection mockConnection;
53+
54+
@BeforeEach
55+
void setUp() {
56+
MockitoAnnotations.openMocks(this);
57+
queryBuilder = new SnowflakeQueryStringBuilder("\"", null);
58+
}
59+
60+
@Test
61+
void testGetFromClauseWithSplit() {
62+
String result = queryBuilder.getFromClauseWithSplit(null, "public", "users", null);
63+
assertEquals(" FROM \"public\".\"users\" ", result);
64+
}
65+
66+
@Test
67+
void testGetPartitionWhereClauses() {
68+
List<String> result = queryBuilder.getPartitionWhereClauses(null);
69+
assertTrue(result.isEmpty());
70+
}
71+
72+
@Test
73+
void testBuildSqlString_NoConstraints() throws SQLException {
74+
Schema tableSchema = new Schema(List.of(new Field("id", new FieldType(true, new ArrowType.Int(32, true), null), null)));
75+
76+
Constraints constraints = mock(Constraints.class);
77+
when(constraints.getLimit()).thenReturn(0L);
78+
79+
String sql = queryBuilder.buildSqlString(mockConnection, null, "public", "users", tableSchema, constraints, null);
80+
assertTrue(sql.contains("SELECT \"id\" FROM \"public\".\"users\" "));
81+
}
82+
83+
@Test
84+
void testBuildSqlString_WithConstraints() throws SQLException {
85+
Schema tableSchema = new Schema(List.of(new Field("id", new FieldType(true, new ArrowType.Int(32, true), null), null)));
5886

87+
Constraints constraints = mock(Constraints.class);
88+
when(constraints.getLimit()).thenReturn(10L);
89+
90+
String sql = queryBuilder.buildSqlString(mockConnection, null, "public", "users", tableSchema, constraints, null);
91+
assertTrue(sql.contains("LIMIT 10"));
92+
}
93+
94+
@Test
95+
void testQuote() {
96+
String result = queryBuilder.quote("users");
97+
assertEquals("\"users\"", result);
98+
}
99+
100+
@Test
101+
void testSingleQuote() {
102+
String result = queryBuilder.singleQuote("O'Reilly");
103+
assertEquals("'O''Reilly'", result);
104+
}
105+
106+
@Test
107+
void testToPredicate_SingleValue() {
108+
List<TypeAndValue> accumulator = new ArrayList<>();
109+
String predicate = queryBuilder.toPredicate("age", "=", 30, new ArrowType.Int(32, true));
110+
assertEquals("age = 30", predicate);
111+
}
112+
113+
@Test
114+
void testGetObjectForWhereClause_Int() {
115+
Object result = queryBuilder.getObjectForWhereClause("age", 42, new ArrowType.Int(32, true));
116+
assertEquals(42L, result);
117+
}
118+
119+
@Test
120+
void testGetObjectForWhereClause_Decimal() {
121+
Object result = queryBuilder.getObjectForWhereClause("price", new BigDecimal("99.99"), new ArrowType.Decimal(10, 2));
122+
assertEquals(new BigDecimal("99.99"), result);
123+
}
124+
125+
@Test
126+
void testGetObjectForWhereClause_Date() {
127+
Object result = queryBuilder.getObjectForWhereClause("date", "2023-03-15T00:00", new ArrowType.Date(DateUnit.DAY));
128+
assertEquals("2023-03-15 00:00:00", result);
129+
}
130+
131+
@Test
132+
void testGetObjectForWhereClause_Timestamp() {
133+
Object result = queryBuilder.getObjectForWhereClause("created_at", 1711929600000L, new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC"));
134+
assertEquals("2024-04-01 00:00:00", result);
135+
}
136+
137+
@Test
138+
void testToPredicateWithUnsupportedType() {
139+
assertThrows(UnsupportedOperationException.class, () ->
140+
queryBuilder.getObjectForWhereClause("unsupported", "value", new ArrowType.Struct())
141+
);
142+
}
143+
}

0 commit comments

Comments
 (0)