eyalgo
9/11/2014 - 8:09 PM

DbCleanupRule.java

public class DbCleanupRule implements TestRule {
	private final DbConnectionManager connection;

	public DbCleanupRule(DbConnectionManager connection) {
		this.connection = connection;
	}

	@Override
	public Statement apply(Statement base, Description description) {
		return new DbCleanupStatement(base, connection);
	}

	private static final class DbCleanupStatement extends Statement {
		private final Statement base;
		private final DbConnectionManager connection;
		
		private DbCleanupStatement(Statement base, DbConnectionManager connection) {
			this.base = base;
			this.connection = connection;
		}

		@Override
		public void evaluate() throws Throwable {
			try {
				cleanDb();
				base.evaluate();
			} finally {
				cleanDb();
			}
		}
		
		private void cleanDb() {
			connection.doTheCleanup();
		}
	}
}