diff --git a/test/DiffSpec.hs b/test/DiffSpec.hs
index e8643412a..ef76bdda5 100644
--- a/test/DiffSpec.hs
+++ b/test/DiffSpec.hs
@@ -4,16 +4,31 @@ module DiffSpec where
 import Category
 import Data.Functor.Both
 import Data.Functor.Listable ()
+import qualified Data.Syntax as Syntax
+import qualified Data.Syntax.Comment as Comment
+import qualified Data.Syntax.Declaration as Declaration
+import qualified Data.Syntax.Statement as Statement
 import Data.Record
-import RWS
+import Data.Union
 import Diff
 import Info
 import Interpreter
-import Syntax
+import RWS
 import Term
 import Test.Hspec
 import Test.Hspec.LeanCheck
 
+type Syntax = Union
+  '[ Comment.Comment
+   , Declaration.Function
+   , Declaration.Method
+   , Statement.If
+   , Syntax.Context
+   , Syntax.Empty
+   , Syntax.Identifier
+   , []
+   ]
+
 spec :: Spec
 spec = parallel $ do
   let decorate = defaultFeatureVectorDecorator (category . termAnnotation)
@@ -22,14 +37,14 @@ spec = parallel $ do
 
   prop "equal terms produce identity diffs" $
     \ a -> let term = decorate (a :: Term Syntax (Record '[Category])) in
-      diffCost (diffSyntaxTerms term term) `shouldBe` 0
+      diffCost (diffTerms term term) `shouldBe` 0
 
   describe "beforeTerm" $ do
     prop "recovers the before term" $
-      \ a b -> let diff = diffSyntaxTerms a b :: Diff Syntax (Record '[Category]) (Record '[Category]) in
+      \ a b -> let diff = diffTerms a b :: Diff Syntax (Record '[Category]) (Record '[Category]) in
         beforeTerm diff `shouldBe` Just a
 
   describe "afterTerm" $ do
     prop "recovers the after term" $
-      \ a b -> let diff = diffSyntaxTerms a b :: Diff Syntax (Record '[Category]) (Record '[Category]) in
+      \ a b -> let diff = diffTerms a b :: Diff Syntax (Record '[Category]) (Record '[Category]) in
         afterTerm diff `shouldBe` Just b