From a67732329d780cab85003dbc48926e64d2ffaee4 Mon Sep 17 00:00:00 2001
From: Konstantine Rybnikov <k-bx@k-bx.com>
Date: Thu, 25 Jul 2013 17:48:29 +0300
Subject: [PATCH] Add yesodSpecWithSiteGenerator method.

This method "extracts" site from `IO site` action every time, instead
of getting it only once. Helpful for flushing database (to recreate
connections) etc.
---
 yesod-test/Yesod/Test.hs | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/yesod-test/Yesod/Test.hs b/yesod-test/Yesod/Test.hs
index 38837967..b7cdead9 100644
--- a/yesod-test/Yesod/Test.hs
+++ b/yesod-test/Yesod/Test.hs
@@ -27,6 +27,7 @@ module Yesod.Test
     ( -- * Declaring and running your test suite
       yesodSpec
     , YesodSpec
+    , yesodSpecWithSiteGenerator
     , YesodExample
     , YesodSpecTree (..)
     , ydescribe
@@ -205,6 +206,28 @@ yesodSpec site yspecs =
             , yedResponse = Nothing
             }
 
+-- | Same as yesodSpec, but instead of taking already built site it
+-- takes an action which produces site for each test.
+yesodSpecWithSiteGenerator :: YesodDispatch site
+                           => IO site
+                           -> YesodSpec site
+                           -> Hspec.Spec
+yesodSpecWithSiteGenerator getSiteAction yspecs =
+    Core.fromSpecList $ map (unYesod getSiteAction) $ execWriter yspecs
+    where
+      unYesod :: YesodDispatch t
+              => IO t -> YesodSpecTree t -> Core.SpecTree
+      unYesod getSiteAction' (YesodSpecGroup x y) = Core.SpecGroup x $ map (unYesod getSiteAction') y
+      unYesod getSiteAction' (YesodSpecItem x y) = Core.it x $ do
+        site <- getSiteAction'
+        app <- toWaiAppPlain site
+        ST.evalStateT y YesodExampleData
+            { yedApp = app
+            , yedSite = site
+            , yedCookies = M.empty
+            , yedResponse = Nothing
+            }
+
 -- | Describe a single test that keeps cookies, and a reference to the last response.
 yit :: String -> YesodExample site () -> YesodSpec site
 yit label example = tell [YesodSpecItem label example]