Problems with afterAll() in ScalaTest in case of Future usage - scala

Question

At the end of my ScalaTest suite I need to do some DB clean up.
The cleanup itself is a Future. The suite does not call the super.afterAll() which leaves some resources used by the suite (like a web browser and db connections) pending.
Here is relevant pieces of code:
override def afterAll():Unit={
var cleanUpsInProgress = true
DB.cleanUpDeletedSegments(db).onComplete{case _ =>
cleanUpsInProgress = false
}
while(cleanUpsInProgress){}
db.close()
aggregatesDB.close()
super.afterAll()
}
and
def cleanUpDeletedSegments(implicit db:ADMPDB):Future[Int]={
db.run{
segments.filter(_.deleted === 1).delete
}
}
I've debugged and scratched my head for a while and got to conclusion it is not even processing the code in the future's onComplete callback. Even when I substitute Slick's db action with stub Future.successfull(1) I still have everything pending and super.afterAll() gets NOT invoked.
Probably I'm doing something stupidly wrong? Could you help?
Note: Also I do think I need to use this ugly var and while loop here because otherwise the main thread gets completed and the framework which initiates the suite running just closes JVM. Maybe I am wrong here so would be great to hear some comments.
--------------------------UPDATE----------------------
The solution by Tyler works. But when I flatMap one more asynch cleanup (which I actually need to do) then the problem is the same again. The code below freezes and does not call super.afterAll:
override def afterAll():Unit={
val cleanUp = DB.cleanUpDeletedSegments(db).flatMap(_ => DB.cleanUpDeletedSegmentGroups(db))
Await.result(cleanUp, 6 seconds)
db.close()
aggregatesDB.close()
super.afterAll()
}
Await.result also does not throw TimeoutException and from what I see neither completed normally. Any ideas?
It works only if I use Await.result sequentially for each future like below:
override def afterAll():Unit={
val cleanUpSegments = DB.cleanUpDeletedSegments(db)
Await.result(cleanUpSegments, 3 seconds)
val cleanUpSegmentGroups = DB.cleanUpDeletedSegmentGroups(db)
Await.result(cleanUpSegmentGroups, 3 seconds)
db.close()
aggregatesDB.close()
super.afterAll()
}

Answer 1

Its probably just easier to await for your Future cleanup to finish:
import scala.concurrent.Await
import scala.concurrent.duration._
override def afterAll() ={
val future = DB.cleanUpDeletedSegments(db)
Await.result(future, 2 minutes)
aggregatesDB.close()
super.afterAll()
}
You can set the timeout to whatever is reasonable

Answer 2

Use solution by #Tyler. Your solution didn't work because you used non-volatile variable cleanupInProgress from multiple threads.