Destination Snowflake add test to avoid duplicated staged data

airbyte-integrations/connectors/destination-snowflake/src/test/java/io/airbyte/integrations/destination/snowflake/SnowflakeDestinationTest.java

@@ -6,12 +6,35 @@
 
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.*;
 
+import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.ImmutableMap;
+import io.airbyte.commons.io.IOs;
 import io.airbyte.commons.jackson.MoreMappers;
+import io.airbyte.commons.json.Jsons;
+import io.airbyte.db.jdbc.JdbcDatabase;
+import io.airbyte.integrations.base.AirbyteMessageConsumer;
+import io.airbyte.integrations.base.Destination;
+import io.airbyte.protocol.models.AirbyteMessage;
+import io.airbyte.protocol.models.AirbyteRecordMessage;
+import io.airbyte.protocol.models.CatalogHelpers;
+import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
+import io.airbyte.protocol.models.DestinationSyncMode;
+import io.airbyte.protocol.models.Field;
+import io.airbyte.protocol.models.JsonSchemaPrimitive;
 import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Test;
 
+import java.nio.file.Path;
+import java.sql.SQLException;
+import java.time.Instant;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
 public class SnowflakeDestinationTest {
 
   private static final ObjectMapper mapper = MoreMappers.initMapper();
@@ -53,4 +76,54 @@ public void useInsertStrategyTest() {
     assertFalse(SnowflakeDestination.isS3Copy(stubConfig));
   }
 
+  @Test
+  public void testCleanupStageOnFailure() throws Exception {
+
+    JdbcDatabase mockDb = mock(JdbcDatabase.class);
+    SnowflakeStagingSqlOperations sqlOperations = mock(SnowflakeStagingSqlOperations.class);
+    final var testMessages = generateTestMessages();
+    final JsonNode config = Jsons.deserialize(IOs.readFile(Path.of("secrets/insert_config.json")));
+
+    AirbyteMessageConsumer airbyteMessageConsumer = new SnowflakeInternalStagingConsumerFactory()
+            .create(Destination::defaultOutputRecordCollector, mockDb,
+            sqlOperations, new SnowflakeSQLNameTransformer(), config, getCatalog());
+    doThrow(SQLException.class).when(sqlOperations).copyIntoTmpTableFromStage(any(),anyString(),anyString(),anyString());
+
+    airbyteMessageConsumer.start();
+    for (AirbyteMessage m : testMessages) {
+        airbyteMessageConsumer.accept(m);
+    }
+
+    try {
+      airbyteMessageConsumer.close();
+    }catch (Exception e){
+      //do nothing cause it's expected behavior
+    }
+
+    verify(sqlOperations, times(1)).cleanUpStage(any(),anyString());
+  }
+
+  private List<AirbyteMessage> generateTestMessages() {
+    return IntStream.range(0, 3)
+            .boxed()
+            .map(i -> new AirbyteMessage()
+                    .withType(AirbyteMessage.Type.RECORD)
+                    .withRecord(new AirbyteRecordMessage()
+                            .withStream("test")
+                            .withNamespace("test_staging")
+                            .withEmittedAt(Instant.now().toEpochMilli())
+                            .withData(Jsons.jsonNode(ImmutableMap.of("id", i, "name", "human " + i)))))
+            .collect(Collectors.toList());
+  }
+
+  ConfiguredAirbyteCatalog getCatalog() {
+    return new ConfiguredAirbyteCatalog().withStreams(List.of(
+            CatalogHelpers.createConfiguredAirbyteStream(
+                            "test",
+                            "test_staging",
+                            Field.of("id", JsonSchemaPrimitive.NUMBER),
+                            Field.of("name", JsonSchemaPrimitive.STRING))
+                    .withDestinationSyncMode(DestinationSyncMode.OVERWRITE)));
+  }
+
 }