Test.Hspec.JSON

JSON.hs

module Test.Hspec.JSON
    ( shouldBeJson
    ) where

import Control.Monad (when)
import Control.Monad.State (StateT, get, modify, runStateT)
import Control.Monad.Writer (Writer, execWriter, tell)
import Data.ByteString.Lazy (ByteString)
import Data.Function (on)
import Data.Monoid ((<>))
import Data.Text (Text)
import Test.Hspec.Expectations.Pretty (Expectation, shouldBe)

import qualified Data.Aeson as Aeson
import qualified Data.ByteString.Lazy.Char8 as Char8
import qualified Data.HashMap.Strict as HashMap
import qualified Data.Text as Text
import qualified Data.Vector as Vector

newtype PrettyJSON = PrettyJSON ByteString

instance Show PrettyJSON where
    show (PrettyJSON s) =
        case Aeson.decode s of
            Just v -> show $ prettyValue v
            Nothing -> "Invalid JSON String: " <> Char8.unpack s

instance Eq PrettyJSON where
    (==) = (==) `on` recode

type Indented = StateT (Bool, Int) (Writer Text)

prettyValue :: Aeson.Value -> Text
prettyValue = runIndented . indentedValue

runIndented :: Indented () -> Text
runIndented i = execWriter $ runStateT i (True, 0)

newline :: Indented ()
newline = do
    tell "\n"
    modify (\(_, i) -> (True, i))

write :: Text -> Indented ()
write t = do
    (beginningOfLine, indentCount) <- get
    modify (\(_, i) -> (False, i))
    when beginningOfLine $ tell $ Text.replicate indentCount "  "
    tell t

indentedValue :: Aeson.Value -> Indented ()
indentedValue (Aeson.Object o) = do
    write "{"
    indent
    newline
    indentedItems (uncurry indentedKeyValue) $ HashMap.toList o
    unindent
    write "}"
indentedValue (Aeson.Array xs) = do
    write "["
    indent
    newline
    indentedItems indentedValue $ Vector.toList xs
    unindent
    write "]"
indentedValue (Aeson.String s) = write ("\"" <> s <> "\"")
indentedValue (Aeson.Number n) = write (Text.pack $ show n)
indentedValue (Aeson.Bool True) = write "true"
indentedValue (Aeson.Bool False) = write "false"
indentedValue Aeson.Null = write "null"

indentedItems :: (a -> Indented ()) -> [a] -> Indented ()
indentedItems _ [] = return ()
indentedItems f (x:[]) = f x >> newline
indentedItems f (x:xs) = f x >> write "," >> newline >> indentedItems f xs

indent :: Indented ()
indent = modify (fmap (+1))

unindent :: Indented ()
unindent = modify (fmap (subtract 1))

indentedKeyValue :: Text -> Aeson.Value -> Indented ()
indentedKeyValue t v = write (t <> ": ") >> indentedValue v

recode :: PrettyJSON -> String
recode (PrettyJSON x) =
    Char8.unpack $ Aeson.encode (Aeson.decode x :: Maybe Aeson.Value)

shouldBeJson :: ByteString -> ByteString -> Expectation
shouldBeJson a b = PrettyJSON a `shouldBe` PrettyJSON b