diff --git a/src/Text/PrettyPrint/Annotated/WL.hs b/src/Text/PrettyPrint/Annotated/WL.hs
--- a/src/Text/PrettyPrint/Annotated/WL.hs
+++ b/src/Text/PrettyPrint/Annotated/WL.hs
@@ -841,11 +841,15 @@
 flatten (Cat x y)       = Cat (flatten x) (flatten y)
 flatten (Nest i x)      = Nest i (flatten x)
 flatten (Union x _)     = flatten x
+flatten (Annotate a x)  = Annotate a (flatten x)
 flatten (Column f)      = Column (flatten . f)
 flatten (Nesting f)     = Nesting (flatten . f)
 flatten (Columns f)     = Columns (flatten . f)
 flatten (Ribbon f)      = Ribbon (flatten . f)
-flatten other           = other                     --Empty,Char,Text
+flatten a@Empty{}       = a
+flatten a@Char{}        = a
+flatten a@Text{}        = a
+flatten a@Line{}        = a
 
 -----------------------------------------------------------
 -- Renderers
@@ -930,7 +934,7 @@
                -> SimpleDoc a)
               -> Float -> Int -> Doc a -> SimpleDoc a
 renderFits nicest rfrac w x
-    = best 0 0 SEmpty (Cons 0 x Nil)
+    = best 0 0 (\_ _ -> SEmpty) (Cons 0 x Nil)
     where
       -- r :: the ribbon width in characters
       r  = max 0 (min w (round (fromIntegral w * rfrac)))
@@ -938,7 +942,7 @@
       -- best :: n = indentation of current line
       --         k = current column
       --        (ie. (k >= n) && (k - n == count of inserted characters)
-      best _ _ z Nil            = z
+      best n k z Nil           = z n k
       best n k z (Cons i d ds) =
         case d of
           Empty         -> best n k z ds
@@ -948,7 +952,8 @@
           FlatAlt l _   -> best n k z (Cons i l ds)
           Cat x' y      -> best n k z (Cons i x' (Cons i y ds))
           Nest j x'     -> let i' = i+j in seq i' (best n k z (Cons i' x' ds))
-          Annotate a d' -> SPushAnn a (best n k (SPopAnn a $ best n k z ds) (Cons i d' Nil))
+          Annotate a d' -> let z' n' k' = SPopAnn a $ best n' k' z ds
+                           in SPushAnn a (best n k z' (Cons i d' Nil))
           Union p q     -> nicest n k w r (best n k z (Cons i p ds))
                                           (best n k z (Cons i q ds))
           Column f      -> best n k z (Cons i (f k) ds)
diff --git a/test/WLPPrintTests.hs b/test/WLPPrintTests.hs
--- a/test/WLPPrintTests.hs
+++ b/test/WLPPrintTests.hs
@@ -1,19 +1,18 @@
 module Main where
 
-import Test.Framework
-import Test.Framework.Providers.HUnit
-import Test.HUnit hiding (Test)
+import Test.Tasty
+import Test.Tasty.HUnit
 import Text.PrettyPrint.Annotated.WL
 
 main :: IO ()
-main = defaultMain [
+main = defaultMain $ testGroup "Tests" [
          testGroup "Tests for each data constructor" conTests
        , testGroup "Tests for some combinators" codeTests
        , testGroup "Tests for the formatting algorithms" formatTests
        , testGroup "Tests for the code examples in the documentation" docTests
        ]
 
-conTests :: [Test]
+conTests :: [TestTree]
 conTests = [
     testCase "Empty tests"   emptyTests
   , testCase "Char tests"    charTests
@@ -30,15 +29,40 @@
   ]
 
 ------------------------------------------------------------
+-- Helper to add annotations
+
+annotateAll :: Doc a -> Doc ()
+annotateAll = Annotate () . go
+  where go Empty          = Empty
+        go (Char x)       = Char x
+        go (Text i s)     = Text i s
+        go Line           = Line
+        go (FlatAlt l r)  = FlatAlt (annotateAll l) (annotateAll r)
+        go (Cat l r)      = Cat (annotateAll l) (annotateAll r)
+        go (Nest i d)     = Nest i (annotateAll d)
+        go (Union l r)    = Union (annotateAll l) (annotateAll r)
+        go (Annotate _ d) = annotateAll d
+        go (Column f)     = Column (annotateAll . f)
+        go (Nesting k)    = Nesting (annotateAll . k)
+        go (Columns k)    = Columns (annotateAll . k)
+        go (Ribbon k)     = Ribbon (annotateAll . k)
+
+------------------------------------------------------------
 -- We test the @Doc@ constructors.
 
 assertPretty :: Int -> String -> String -> Doc a -> Assertion
-assertPretty w desc str doc = assertEqual (desc ++ " (pretty)") str
-                                $ displayS (renderPretty 1.0 w doc) ""
+assertPretty w desc str doc = do
+  assertEqual (desc ++ " (pretty)") str
+    $ displayS (renderPretty 1.0 w doc) ""
+  assertEqual (desc ++ " (pretty, annotated)") str
+    $ displayS (renderPretty 1.0 w $ annotateAll doc) ""
 
 assertSmart :: Int -> String -> String -> Doc a -> Assertion
-assertSmart w desc str doc = assertEqual (desc ++ " (smart)") str
-                                $ displayS (renderSmart w doc) ""
+assertSmart w desc str doc = do
+  assertEqual (desc ++ " (smart)") str
+    $ displayS (renderSmart w doc) ""
+  assertEqual (desc ++ " (smart, annotated)") str
+    $ displayS (renderSmart w $ annotateAll doc) ""
 
 assertRender :: Int -> String -> String -> Doc a -> Assertion
 assertRender w desc str doc = do assertPretty w desc str doc
@@ -63,7 +87,7 @@
                     $ flatten $ flatAlt (text "x") (text "y")
 
 catTests :: Assertion
-catTests = do assertRender 80 "Cat test 1" "some code"
+catTests = assertRender 80 "Cat test 1" "some code"
                 $ text "some" <> space <> text "code"
 
 nestTests :: Assertion
@@ -79,25 +103,25 @@
                   $ text "foo" </> text "bar"
 
 columnTests :: Assertion
-columnTests = do assertRender 80 "Column test 1" "foo 4"
-                   $ text "foo" <+> (column $ \c -> pretty c)
+columnTests = assertRender 80 "Column test 1" "foo 4"
+                 $ text "foo" <+> column pretty
 
 nestingTests :: Assertion
-nestingTests = do assertRender 80 "Nesting test 1" "foo 2"
-                    $ text "foo" <+> (nest 2 $ nesting $ \n -> pretty n)
+nestingTests =  assertRender 80 "Nesting test 1" "foo 2"
+                  $ text "foo" <+> nest 2 (nesting pretty)
 
 columnsTests :: Assertion
-columnsTests = do assertRender 21 "Columns test 1" "foo 21"
-                    $ text "foo" <+> (nest 2 $ columns $ \w -> pretty w)
+columnsTests = assertRender 21 "Columns test 1" "foo 21"
+                  $ text "foo" <+> nest 2 (columns pretty)
 
 ribbonTests :: Assertion
-ribbonTests = do assertEqual "Ribbon test 1" "foo 40"
-                   $ show (text "foo" <+> (ribbon $ \r -> pretty r))
+ribbonTests = assertEqual "Ribbon test 1" "foo 40"
+                 $ show (text "foo" <+> ribbon pretty)
 
 ------------------------------------------------------------
 -- We test some combinators.
 
-codeTests :: [Test]
+codeTests :: [TestTree]
 codeTests = [
     testCase "@list@ tests"   listTests
   , testCase "@tupled@ tests" tupledTests
@@ -119,7 +143,7 @@
 -- We test some corner cases of the formatting algorithms on a prototypical
 -- syntax for a scripting language (e.g. Python).
 
-formatTests :: [Test]
+formatTests :: [TestTree]
 formatTests = [
     testCase "@renderPretty@ test" renderPrettyTest
   , testCase "@renderSmart@ test"  renderSmartTest
@@ -216,7 +240,7 @@
 ------------------------------------------------------------
 -- We test the code examples in the haddock comments.
 
-docTests :: [Test]
+docTests :: [TestTree]
 docTests = [
     testCase "@fill@ test" fillTest
   , testCase "@fillBreak@ test" fillBreakTest
@@ -230,7 +254,7 @@
         , ("linebreak","Doc e") ]
 
 fillTest :: Assertion
-fillTest = do assertRender 80 "@fill@ test 1" (concat [
+fillTest = assertRender 80 "@fill@ test 1" (concat [
                   "let empty  :: Doc e\n"
                 , "    nest   :: Int -> Doc e -> Doc e\n"
                 , "    linebreak :: Doc e" ])
@@ -239,7 +263,7 @@
                       fill 6 (text name) <+> text "::" <+> text tp
 
 fillBreakTest :: Assertion
-fillBreakTest = do assertRender 80 "@fillBreak@ test 1" (concat [
+fillBreakTest = assertRender 80 "@fillBreak@ test 1" (concat [
                        "let empty  :: Doc e\n"
                      , "    nest   :: Int -> Doc e -> Doc e\n"
                      , "    linebreak\n"
@@ -249,7 +273,7 @@
                         fillBreak 6 (text name) <+> (text "::" <+> text tp)
 
 hangTest :: Assertion
-hangTest = do assertRender 20 "@hang@ test 1" (concat [
+hangTest = assertRender 20 "@hang@ test 1" (concat [
                   "the hang combinator\n"
                 , "    indents these\n"
                 , "    words !" ])
@@ -257,7 +281,7 @@
                 $ words "the hang combinator indents these words !"
 
 alignTest :: Assertion
-alignTest = do assertRender 20 "@align@ test 1" (concat [
+alignTest = assertRender 20 "@align@ test 1" (concat [
                    "hi nice\n"
                  , "   world" ])
                  $ text "hi" <+> (text "nice" $$ text "world")
diff --git a/wl-pprint-annotated.cabal b/wl-pprint-annotated.cabal
--- a/wl-pprint-annotated.cabal
+++ b/wl-pprint-annotated.cabal
@@ -1,10 +1,10 @@
--- This file has been generated from package.yaml by hpack version 0.15.0.
+-- This file has been generated from package.yaml by hpack version 0.17.0.
 --
 -- see: https://github.com/sol/hpack
 
 name:          wl-pprint-annotated
 category:      Text
-version:       0.0.1.4
+version:       0.1.0.0
 synopsis:      Wadler/Leijen pretty printer with annotations and slightly modernized API
 cabal-version: >= 1.10
 license:       BSD3
@@ -53,9 +53,8 @@
     text       >= 0.11 && < 1.3,
     deepseq    >= 1.4  && < 1.6,
     wl-pprint-annotated,
-    HUnit,
-    test-framework,
-    test-framework-hunit
+    tasty,
+    tasty-hunit
   if impl(ghc < 8.0)
     build-depends:
       semigroups >= 0.9 && < 1