Mocking IO Actions: getArgs and putStrLn - unit-testing

Question

I'm trying to test a small function (or rather, IO Action) that takes a command line argument and outputs it to the screen. My original (untestable) function is:
-- In Library.hs
module Library where
import System.Environment (getArgs)
run :: IO ()
run = do
args <- getArgs
putStrLn $ head args
After looking at this answer about mocking, I have come up with a way to mock getArgs and putStrLn by using a type class constrained type. So the above function becomes:
-- In Library.hs
module Library where
class Monad m => SystemMonad m where
getArgs :: m [String]
putStrLn :: String -> m ()
instance SystemMonad IO where
getArgs = System.Environment.getArgs
putStrLn = Prelude.putStrLn
run :: SystemMonad m => m ()
run = do
args <- Library.getArgs
Library.putStrLn $ head args
This Library., Prelude. and System.Environment. are to avoid compiler complaints of Ambigious Occurence. My test file looks like the following.
-- In LibrarySpec.hs
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
import Library
import Test.Hspec
import Control.Monad.State
data MockArgsAndResult = MockArgsAndResult [String] String
deriving(Eq, Show)
instance SystemMonad (State MockArgsAndResult) where
getArgs = do
MockArgsAndResult args _ <- get
return args
putStrLn string = do
MockArgsAndResult args _ <- get
put $ MockArgsAndResult args string
return ()
main :: IO ()
main = hspec $ do
describe "run" $ do
it "passes the first command line argument to putStrLn" $ do
(execState run (MockArgsAndResult ["first", "second"] "")) `shouldBe` (MockArgsAndResult ["first", "second"] "first")
I'm using a State monad that effectively contains 2 fields.
A list for the command line arguments where the mock getArgs reads from
A string that the mock putStrLn puts what was passed to it.
The above code works and seems to test what I want it to test. However, I'm wondering if there is some better / cleaner / more idiomatic way of testing this. For one thing, I'm using the same state to both put stuff into the test (my fake command line arguments), and then get stuff out of it (what was passed to putStrLn.
Is there a better way of doing what I'm doing? I'm more familiar with mocking in a Javascript environment, and my knowledge of Haskell is pretty basic (I arrived at the above solution by a fair bit of trial and error, rather than actual understanding)

Answer 1

The better way is to avoid needing to provide mock versions of getArgs and putStrLn by separating out the heart of the computation into a pure function.
Consider this example:
main = do
args <- getArgs
let n = length $ filter (\w -> length w < 5) args
putStrLn $ "Number of small words: " ++ show n
One could say that the heart of the computation is counting the number of small words which is a pure function of type [String] -> Int. This suggest that we should refactor the program like this:
main = do
args <- getArgs
let n = countSmallWords args
putStrLn $ "Number of small words: " ++ show n
countSmallWords :: [String] -> Int
countSmallWords ws = ...
Now we just test countSmallWords, and this is easy because it is pure function.