Очистка БД и сброс autoincrement полей в тестах junit

Очистка БД и сброс autoincrement полей в тестах junit

@Slf4j
@ActiveProfiles("test")
@SpringBootTest
@AutoConfigureMockMvc
public abstract class BaseITest {
	@AfterEach
	public void down(@Autowired JdbcTemplate jdbcTemplate) {
		cleanTables(jdbcTemplate);
		cleanSequences(jdbcTemplate, "PUBLIC");
		cleanIdentities(jdbcTemplate, "PUBLIC");
	}

	private void cleanSequences(JdbcTemplate jdbcTemplate, String schema) {
		List<String> sequences = jdbcTemplate.queryForList("SELECT * FROM INFORMATION_SCHEMA.SEQUENCES")
				.stream()
				.filter(stringObjectMap -> stringObjectMap.get("SEQUENCE_SCHEMA").equals(schema))
				.filter(stringObjectMap -> !stringObjectMap.get("CURRENT_VALUE").equals(0L))
				.map(stringObjectMap -> (String) stringObjectMap.get("SEQUENCE_NAME"))
				.collect(Collectors.toList());
		for (String s : sequences) {
			String query = "alter sequence " + s + " restart with 1;";
			jdbcTemplate.update(query);
		}
		log.debug("Clean sequences: " + StringUtils.trimTrailingCharacter(String.join(" ,", sequences), ','));
	}

	private void cleanIdentities(JdbcTemplate jdbcTemplate, String schema) {

		List<String[]> identities = jdbcTemplate.queryForList("SELECT * FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA = '" + schema + "' AND IS_IDENTITY = 'YES'")
				.stream()
				.map(stringObjectMap -> new String[] {(String) stringObjectMap.get("TABLE_NAME"), (String) stringObjectMap.get("COLUMN_NAME")})
				.collect(Collectors.toList());
		for (String[] i : identities) {
			String query = "ALTER TABLE " + i[0] + " ALTER COLUMN " + i[1] + " RESTART WITH 1";
			jdbcTemplate.update(query);
		}
		log.debug("Clean identities: " + StringUtils.trimTrailingCharacter(identities.stream().map(strings -> strings[0] + "." + strings[1]).collect(Collectors.joining(" ,")), ','));
	}

	private void cleanTables(JdbcTemplate jdbcTemplate) {
		StringBuilder toLog = new StringBuilder();
		List<String> querys = jdbcTemplate
				.queryForList("show tables from public")
				.stream()
				.map(stringObjectMap -> stringObjectMap.get("TABLE_NAME"))
				.filter(o -> !"DATABASECHANGELOG".equals(o))
				.filter(o -> !"DATABASECHANGELOGLOCK".equals(o))
				.map(Object::toString)
				.collect(Collectors.toList());
		while (!querys.isEmpty()) {
			querys.removeIf(s -> {
				int rowCount = JdbcTestUtils.countRowsInTable(jdbcTemplate, s);
				if (rowCount > 0) {
					String query = "DELETE FROM " + s + ";";
					try {
						jdbcTemplate.update(query);
						toLog.append(" " + s + ",");
						return true;
					} catch (Exception e) {
						return false;
					}
				}else{
					return true;
				}
			});
		}
		if (toLog.toString().length() > 0) {
			log.debug("Clean tables:" + StringUtils.trimTrailingCharacter(toLog.toString(), ','));
		}
	}
}
(Просмотрено 113 раз, 1 раз за сегодня)

Добавить комментарий

Ваш адрес email не будет опубликован. Обязательные поля помечены *