tasty-discover.hs

import System.Environment (getArgs)
import System.IO (hPutStrLn, stderr)
import System.Exit (exitFailure)
import Data.List (find, isPrefixOf, isSuffixOf, nub, intersperse)
import Data.Maybe (fromJust, mapMaybe)
import System.Directory (listDirectory, doesDirectoryExist)
import Data.Traversable (for)
import System.FilePath ((</>), takeDirectory, pathSeparator, dropExtension)
import Data.Monoid (Endo(..))
import Data.Foldable (fold)

data Generator = Generator
  { genPrefix :: String
  , genImport :: Maybe String
  , genTree, genSetup :: Test -> ShowS }

data Test = Test { testModule, testFunction :: String }

str :: String -> ShowS
str = (++)

sp, nl :: ShowS
sp = (' ':)
nl = ('\n':)

tr :: Char -> Char -> String -> String
tr a b = map $ \c -> if c == a then b else c

name, fn, setup :: Test -> ShowS
name = shows . tr '_' ' ' . tail . dropWhile (/= '_') . testFunction
fn t = str (testModule t) . ('.':) . str (testFunction t)
setup t = str "setup_" . str (tr '.' '_' $ testModule t) . ('_':) . str (testFunction t)

generators :: [Generator]
generators =
  [ Generator { genPrefix = "prop_"
              , genImport = Just "Test.Tasty.QuickCheck"
              , genSetup = const id
              , genTree = \t -> str "Test.Tasty.QuickCheck.testProperty " . name t . sp . fn t }
  , Generator { genPrefix = "case_"
              , genImport = Just "Test.Tasty.HUnit"
              , genSetup = const id
              , genTree = \t -> str "Test.Tasty.HUnit.testCase " . name t . sp . fn t }
  , Generator { genPrefix = "spec_"
              , genImport = Just "Test.Tasty.Hspec"
              , genSetup = \t -> str "  " . setup t . str " <- Test.Tasty.Hspec.testSpec " . name t . sp . fn t . nl
              , genTree = setup }
  , Generator { genPrefix = "gen_"
              , genImport = Nothing
              , genSetup = \t -> str "  " . setup t . str " <- " . fn t . nl
              , genTree = \t -> str "Test.Tasty.testGroup " . name t . sp . setup t }
  , Generator { genPrefix = "group_"
              , genImport = Nothing
              , genSetup = const id
              , genTree = \t -> str "Test.Tasty.testGroup " . name t . sp . fn t }
  , Generator { genPrefix = "test_"
              , genImport = Nothing
              , genSetup = const id
              , genTree = fn } ]

testFileSuffixes :: [String]
testFileSuffixes = (++) <$> ["Spec", "Test"] <*> [".lhs", ".hs"]

getGenerator :: Test -> Generator
getGenerator t = fromJust $ find ((`isPrefixOf` testFunction t) . genPrefix) generators

showImports :: [Test] -> ShowS
showImports ts = foldEndo $ map (\m -> str "import qualified " . str m . nl) $ nub $
  map testModule ts ++ mapMaybe (genImport . getGenerator) ts

showSetup :: Test -> ShowS
showSetup t = genSetup (getGenerator t) t

showTree :: Test -> ShowS
showTree t = genTree (getGenerator t) t

foldEndo :: (Functor f, Foldable f) => f (a -> a) -> (a -> a)
foldEndo = appEndo . fold . fmap Endo

showModule :: FilePath -> [Test] -> ShowS
showModule src t =
    str "{-# LINE 1 " . shows src . str " #-}\n"
  . str "module Main where\n\n"
  . str "import Prelude\n"
  . str "import qualified Test.Tasty\n"
  . showImports t
  . str "\nmain :: IO ()\n"
  . str "main = do\n"
  . foldEndo (map showSetup t)
  . str "  Test.Tasty.defaultMain $ Test.Tasty.testGroup " . shows src
  . str "\n    [ "
  . foldEndo (intersperse (str "\n    , ") $ map showTree t)
  . str " ]\n"

filesBySuffix :: FilePath -> [String] -> IO [FilePath]
filesBySuffix dir suffixes = do
  entries <- listDirectory dir
  found <- for entries $ \entry -> do
    let dir' = dir </> entry
    exists <- doesDirectoryExist dir'
    if exists then map (entry </>) <$> filesBySuffix dir' suffixes else pure []
  pure $ filter (\x -> any (`isSuffixOf` x) suffixes) entries ++ concat found

findTests :: FilePath -> IO [Test]
findTests src = do
  let dir = takeDirectory src
  files <- filesBySuffix dir testFileSuffixes
  concat <$> traverse (extractTests dir) files

mkTest :: FilePath -> String -> Test
mkTest = Test . tr pathSeparator '.' . dropExtension

extractTests :: FilePath -> FilePath -> IO [Test]
extractTests dir file =
    map (mkTest file) . nub
  . filter (\n -> any ((`isPrefixOf` n) . genPrefix) generators)
  . map fst . concatMap lex . lines <$> readFile (dir </> file)

main :: IO ()
main = do
  args <- getArgs
  case args of
    src : _ : dst : _ -> do
      tests <- findTests src
      writeFile dst $ showModule src tests ""
    _ -> do
      hPutStrLn stderr "tasty-discover: Expected source and destination arguments"
      exitFailure