2 files changed, 34 insertions, 23 deletions
diff --git a/haddock-test/src/Test/Haddock/Xhtml.hs b/haddock-test/src/Test/Haddock/Xhtml.hs
index 35f5910a..b6941496 100644
--- a/haddock-test/src/Test/Haddock/Xhtml.hs
+++ b/haddock-test/src/Test/Haddock/Xhtml.hs
@@ -2,47 +2,55 @@
 {-# LANGUAGE StandaloneDeriving #-}
 
 
-module Test.Haddock.Xhtml where
+module Test.Haddock.Xhtml
+    ( Xhtml(..)
+    , parseXhtml, dumpXhtml
+    , stripLinks, stripFooter
+    ) where
 
 
-import Control.Monad
-
 import Data.Generics.Aliases
 import Data.Generics.Schemes
 
 import Text.XML.Light
 
 
-deriving instance Eq Content
+newtype Xhtml = Xhtml
+    { xhtmlElement :: Element
+    } deriving Eq
+
+
+-- TODO: Find a way to avoid warning about orphan instances.
 deriving instance Eq Element
+deriving instance Eq Content
 deriving instance Eq CData
 
 
-readXml :: FilePath -> IO (Maybe Element)
-readXml = liftM parseXMLDoc . readFile
+parseXhtml :: String -> Maybe Xhtml
+parseXhtml = fmap Xhtml . parseXMLDoc
 
 
-strip :: Element -> Element
-strip = stripFooter . stripLinks
+dumpXhtml :: Xhtml -> String
+dumpXhtml = ppElement . xhtmlElement
 
 
-stripLinks :: Element -> Element
+stripLinks :: Xhtml -> Xhtml
 stripLinks =
-    everywhere (mkT unlink)
+    Xhtml . everywhere (mkT unlink) . xhtmlElement
   where
     unlink attr@(Attr { attrKey = key })
         | qName key == "href" = attr { attrVal = "#" }
         | otherwise = attr
 
 
-stripFooter :: Element -> Element
+stripFooter :: Xhtml -> Xhtml
 stripFooter =
-    everywhere (mkT defoot)
+    Xhtml . everywhere (mkT defoot) . xhtmlElement
   where
-    defoot elem
-        | isFooter elem = elem { elContent = [] }
-        | otherwise = elem
-    isFooter elem = any isFooterAttr $ elAttribs elem
+    defoot el
+        | isFooter el = el { elContent = [] }
+        | otherwise = el
+    isFooter el = any isFooterAttr $ elAttribs el
     isFooterAttr (Attr { .. }) = and
         [ qName attrKey == "id"
         , attrVal == "footer"
diff --git a/html-test/run.hs b/html-test/run.hs
index 2758bf56..ab007f57 100755
--- a/html-test/run.hs
+++ b/html-test/run.hs
@@ -7,13 +7,11 @@ import System.FilePath
 import Test.Haddock
 import Test.Haddock.Xhtml
 
-import qualified Text.XML.Light as Xml
 
-
-checkConfig :: CheckConfig Xml.Element
+checkConfig :: CheckConfig Xhtml
 checkConfig = CheckConfig
-    { ccfgRead = \_ input -> strip <$> Xml.parseXMLDoc input
-    , ccfgDump = Xml.ppElement
+    { ccfgRead = \mdl input -> stripIfRequired mdl <$> parseXhtml input
+    , ccfgDump = dumpXhtml
     , ccfgEqual = (==)
     }
 
@@ -26,8 +24,13 @@ main :: IO ()
 main = runAndCheck =<< parseArgs checkConfig dirConfig =<< getArgs
 
 
--- *** OLD TEST RUNNER UTILITY FUNCTIONS ***
--- These are considered bad and should be replaced as soon as possible.
+stripIfRequired :: String -> Xhtml -> Xhtml
+stripIfRequired mdl =
+    stripLinks' . stripFooter
+  where
+    stripLinks'
+        | mdl `elem` preserveLinksModules = id
+        | otherwise = stripFooter
 
 
 -- | List of modules in which we don't 'stripLinks'