2 files changed, 41 insertions, 13 deletions

diff --git a/README.md b/README.md
index e254c58..02ed812 100644
--- a/README.md
+++ b/README.md
@@ -1,11 +1,18 @@
 # squeal-hspec
 
-Helpers for creating database tests with hspec and squeal, inspired by Jonathan Fischoff's
+Helpers for creating database tests with hspec and squeal, inspired by Jonathan
+Fischoff's
 [hspec-pg-transact](http://hackage.haskell.org/package/hspec-pg-transact).
 
-This uses @tmp-postgres@ to automatically and connect to a temporary instance of postgres on a random port.
+This uses @tmp-postgres@ to automatically and connect to a temporary instance of
+postgres on a random port.
 
 Current version is done to operate with Squeal 0.4.0.0.
 
-`describeDB` lets you initate a series of `itDB` specs which will operate in the same context.
-It takes migrations to run and a series of fixtures to fill the database.
+`describeDB` lets you initate a series of `itDB` specs which will operate in the
+same context. It takes migrations to run and a series of fixtures to fill the
+database.
+
+Setting the env var `TEST_DB_CONNECTION_STRING` will let you use a non-temporary
+database if need be. Tear down will run the `migrateDown` to make sure the test
+db doesn't get filled with invalid data.
diff --git a/src/Squeal/PostgreSQL/Hspec.hs b/src/Squeal/PostgreSQL/Hspec.hs
index 3cd886b..d065048 100644
--- a/src/Squeal/PostgreSQL/Hspec.hs
+++ b/src/Squeal/PostgreSQL/Hspec.hs
@@ -13,6 +13,7 @@ The libary also provides a few other functions for more fine grained control ove
 {-# LANGUAGE MonoLocalBinds   #-}
 {-# LANGUAGE RankNTypes       #-}
 {-# LANGUAGE RecordWildCards  #-}
+{-# LANGUAGE TupleSections    #-}
 {-# LANGUAGE TypeInType       #-}
 {-# LANGUAGE TypeOperators    #-}
 module Squeal.PostgreSQL.Hspec
@@ -21,33 +22,49 @@ where
 import           Control.Exception
 import           Control.Monad
 import           Control.Monad.Base     (liftBase)
+import           Data.ByteString        (ByteString)
 import qualified Data.ByteString.Char8  as BSC
 import qualified Database.Postgres.Temp as Temp
 import           Generics.SOP           (K)
 import           Squeal.PostgreSQL
 import           Squeal.PostgreSQL.Pool
+import           System.Environment     (lookupEnv)
 import           Test.Hspec
 
 data TestDB a = TestDB
-  { tempDB :: Temp.DB
+  { tempDB           :: Maybe Temp.DB
   -- ^ Handle for temporary @postgres@ process
-  , pool   :: Pool a
+  , pool             :: Pool a
   -- ^ Pool of 50 connections to the temporary @postgres@
+  , connectionString :: ByteString
   }
 
 type Fixtures schema = (Pool (K Connection schema) -> IO ())
 type Actions schema a = PoolPQ schema IO a
 type SquealContext schema = TestDB (K Connection schema)
 
+testDBEnv :: String
+testDBEnv = "TEST_DB_CONNECTION_STRING"
+
+getOrCreateConnectionString :: IO (ByteString, Maybe Temp.DB)
+getOrCreateConnectionString = do
+  hasConnectionString <- lookupEnv testDBEnv
+  maybe createTempDB (pure . (, Nothing) . BSC.pack) hasConnectionString
+
+createTempDB :: IO (ByteString, Maybe Temp.DB)
+createTempDB = do
+  tempDB <- either throwIO return =<< Temp.startAndLogToTmp []
+  let connectionString = BSC.pack (Temp.connectionString tempDB)
+  pure (connectionString, Just tempDB)
+
 -- | Start a temporary @postgres@ process and create a pool of connections to it
 setupDB
   :: AlignedList (Migration (Terminally PQ IO)) schema0 schema
   -> Fixtures schema
   -> IO (SquealContext schema)
 setupDB migration fixtures = do
-  tempDB <- either throwIO return =<< Temp.startAndLogToTmp []
-  let connectionString = BSC.pack (Temp.connectionString tempDB)
-  putStrLn $ Temp.connectionString tempDB
+  (connectionString, tempDB) <- getOrCreateConnectionString
+  BSC.putStrLn connectionString
   let singleStripe = 1
       keepConnectionForOneHour = 3600
       poolSizeOfFifty = 50
@@ -61,10 +78,14 @@ setupDB migration fixtures = do
   pure TestDB {..}
 
 -- | Drop all the connections and shutdown the @postgres@ process
-teardownDB :: TestDB a -> IO ()
-teardownDB TestDB {..} = do
+teardownDB
+  :: AlignedList (Migration (Terminally PQ IO)) schema0 schema
+  -> TestDB a
+  -> IO ()
+teardownDB migration TestDB {..} = do
+  withConnection connectionString (migrateDown migration)
   destroyAllResources pool
-  void $ Temp.stop tempDB
+  maybe (pure ()) (void . Temp.stop) tempDB
 
 -- | Run an 'IO' action with a connection from the pool
 withPool :: TestDB (K Connection schema) -> Actions schema a -> IO a
@@ -104,4 +125,4 @@ describeDB
   -> SpecWith (SquealContext schema)
   -> Spec
 describeDB migrate fixture str =
-  beforeAll (setupDB migrate fixture) . afterAll teardownDB . describe str
+  beforeAll (setupDB migrate fixture) . afterAll (teardownDB migrate) . describe str