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

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