@@ -65,22 +65,24 @@ public void testGetSchemaInvalidJson() {
6565
6666 @ Test
6767 public void testGetSchemaFromSnowflakeUnknownType () throws IOException {
68+ String importQuery = "SELECT * FROM someTable" ;
6869 MockFailureCollector collector = new MockFailureCollector (MOCK_STAGE );
6970 SnowflakeSourceAccessor snowflakeAccessor = Mockito .mock (SnowflakeSourceAccessor .class );
7071
7172 List <SnowflakeFieldDescriptor > sample = new ArrayList <>();
7273 sample .add (new SnowflakeFieldDescriptor ("field1" , -1000 , false ));
7374
74- Mockito .when (snowflakeAccessor .describeQuery (null )).thenReturn (sample );
75+ Mockito .when (snowflakeAccessor .describeQuery (importQuery )).thenReturn (sample );
7576
76- SchemaHelper .getSchema (snowflakeAccessor , null , collector , null );
77+ SchemaHelper .getSchema (snowflakeAccessor , null , collector , importQuery );
7778
7879 ValidationAssertions .assertValidationFailed (
7980 collector , Collections .singletonList (SnowflakeBatchSourceConfig .PROPERTY_SCHEMA ));
8081 }
8182
8283 @ Test
8384 public void testGetSchemaFromSnowflake () throws IOException {
85+ String importQuery = "SELECT * FROM someTable" ;
8486 MockFailureCollector collector = new MockFailureCollector (MOCK_STAGE );
8587 SnowflakeSourceAccessor snowflakeAccessor = Mockito .mock (SnowflakeSourceAccessor .class );
8688
@@ -142,9 +144,9 @@ public void testGetSchemaFromSnowflake() throws IOException {
142144 Schema .Field .of ("field134" , Schema .nullableOf (Schema .of (Schema .LogicalType .TIMESTAMP_MICROS )))
143145 );
144146
145- Mockito .when (snowflakeAccessor .describeQuery (null )).thenReturn (sample );
147+ Mockito .when (snowflakeAccessor .describeQuery (importQuery )).thenReturn (sample );
146148
147- Schema actual = SchemaHelper .getSchema (snowflakeAccessor , null , collector , null );
149+ Schema actual = SchemaHelper .getSchema (snowflakeAccessor , null , collector , importQuery );
148150
149151 Assert .assertTrue (collector .getValidationFailures ().isEmpty ());
150152 Assert .assertEquals (expected , actual );
@@ -182,4 +184,21 @@ public void testGetSchemaWhenMacroIsEnabledSchemaIsNull() {
182184 Assert .assertNull (actual );
183185
184186 }
187+
188+ @ Test
189+ public void testGetSchemaManuallyUpdatedTheSchema () {
190+ Schema expected = Schema .recordOf ("test" ,
191+ Schema .Field .of ("test_field" , Schema .nullableOf (Schema .of (Schema .Type .LONG )))
192+ );
193+
194+ SnowflakeBatchSourceConfig mockConfig = Mockito .mock (SnowflakeBatchSourceConfig .class );
195+ Mockito .when (mockConfig .canConnect ()).thenReturn (false );
196+ Mockito .when (mockConfig .getSchema ()).thenReturn (expected .toString ());
197+
198+ MockFailureCollector collector = new MockFailureCollector (MOCK_STAGE );
199+ Schema actual = SchemaHelper .getSchema (mockConfig , collector );
200+
201+ Assert .assertTrue (collector .getValidationFailures ().isEmpty ());
202+ Assert .assertEquals (expected , actual );
203+ }
185204}
0 commit comments