diff --git a/hydra-node/test/Hydra/Chain/Direct/TxTraceSpec.hs b/hydra-node/test/Hydra/Chain/Direct/TxTraceSpec.hs
index 5459260c659..33ac0d79a3f 100644
--- a/hydra-node/test/Hydra/Chain/Direct/TxTraceSpec.hs
+++ b/hydra-node/test/Hydra/Chain/Direct/TxTraceSpec.hs
@@ -14,15 +14,11 @@ import Test.Hydra.Prelude
 
 import Cardano.Api.UTxO (UTxO)
 import Cardano.Api.UTxO qualified as UTxO
-import Data.List qualified as List
 import Data.Map.Strict qualified as Map
 import Data.Time.Clock.POSIX (posixSecondsToUTCTime)
 import Hydra.Cardano.Api (
-  Coin,
   SlotNo (..),
   mkTxOutDatumInline,
-  modifyTxOutValue,
-  renderUTxO,
   selectLovelace,
   throwError,
   txOutAddress,
@@ -122,43 +118,11 @@ prop_runActions actions =
 
 data Model = Model
   { headState :: State
-  , latestSnapshot :: ModelSnapshot
+  , latestSnapshot :: SnapshotNumber
   , alreadyContested :: [Actor]
   }
   deriving (Show)
 
--- | A snapshot that may have pending decommits (= augmented).
-data ModelSnapshot = Normal SnapshotNumber | Augmented SnapshotNumber
-  deriving (Show, Eq)
-
-snapshotNumber :: ModelSnapshot -> SnapshotNumber
-snapshotNumber = \case
-  Normal n -> n
-  Augmented n -> n
-
-isAugmented :: ModelSnapshot -> Bool
-isAugmented = \case
-  Normal{} -> False
-  Augmented{} -> True
-
-instance Ord ModelSnapshot where
-  compare a b = compare (snapshotNumber a) (snapshotNumber b)
-
-instance Num ModelSnapshot where
-  a + b = Normal $ snapshotNumber a + snapshotNumber b
-  a - b = Normal $ snapshotNumber a - snapshotNumber b
-  a * b = Normal $ snapshotNumber a * snapshotNumber b
-  abs = Normal . abs . snapshotNumber
-  signum = Normal . signum . snapshotNumber
-  fromInteger = Normal . fromInteger
-
-instance Arbitrary ModelSnapshot where
-  arbitrary = oneof [Normal <$> arbitrary, Augmented <$> arbitrary]
-
-  shrink = \case
-    Normal n -> Normal <$> shrink n
-    Augmented n -> Augmented <$> shrink n
-
 data State
   = Open
   | Closed
@@ -177,10 +141,10 @@ data TxResult = TxResult
 
 instance StateModel Model where
   data Action Model a where
-    Decrement :: {actor :: Actor, snapshot :: ModelSnapshot} -> Action Model TxResult
-    Close :: {actor :: Actor, snapshot :: ModelSnapshot} -> Action Model TxResult
-    Contest :: {actor :: Actor, snapshot :: ModelSnapshot} -> Action Model TxResult
-    Fanout :: {snapshot :: ModelSnapshot} -> Action Model TxResult
+    Decrement :: {actor :: Actor, snapshot :: SnapshotNumber} -> Action Model TxResult
+    Close :: {actor :: Actor, snapshot :: SnapshotNumber} -> Action Model TxResult
+    Contest :: {actor :: Actor, snapshot :: SnapshotNumber} -> Action Model TxResult
+    Fanout :: {snapshot :: SnapshotNumber} -> Action Model TxResult
     -- \| Helper action to identify the terminal state 'Final' and shorten
     -- traces using the 'precondition'.
     Stop :: Action Model ()
@@ -203,7 +167,7 @@ instance StateModel Model where
               pure $ Some $ Close{actor, snapshot}
           , do
               actor <- elements allActors
-              snapshot <- (latestSnapshot + 1) `orSometimes` (Augmented <$> arbitrary)
+              snapshot <- (latestSnapshot + 1) `orSometimes` arbitrary
               pure $ Some Decrement{actor, snapshot}
           ]
       Closed{} ->
@@ -225,10 +189,8 @@ instance StateModel Model where
       -- TODO: assert what to decrement still there
       headState == Open
         && snapshot > latestSnapshot
-        && isAugmented snapshot
     Close{snapshot} ->
-      snapshot /= Augmented 0 -- TODO: don't generate this one
-        && headState == Open
+      headState == Open
         && snapshot >= latestSnapshot
     Contest{actor, snapshot} ->
       headState == Closed
@@ -252,7 +214,7 @@ instance StateModel Model where
     Fanout{snapshot} ->
       headState == Closed -- TODO: gracefully fail in perform instead?
       -- TODO: why can't we have this condition too? It causes CannotFindHeadOutput... errors
-      -- && snapshot /= latestSnapshot
+        && snapshot /= latestSnapshot
     _ -> False
 
   nextState :: Model -> Action Model a -> Var a -> Model
@@ -407,11 +369,11 @@ getValidationError tx utxo =
 allActors :: [Actor]
 allActors = [Alice, Bob, Carol]
 
--- | A "random" UTxO distribution for a given 'ModelSnapshot'. This always
+-- | A "random" UTxO distribution for a given 'SnapshotNumber'. This always
 -- contains one UTxO for alice, bob, and carol.
-snapshotUTxO :: ModelSnapshot -> UTxO
+snapshotUTxO :: SnapshotNumber -> UTxO
 snapshotUTxO snapshot =
-  (`generateWith` fromIntegral (snapshotNumber snapshot)) . resize 1 $ do
+  (`generateWith` fromIntegral snapshot) . resize 1 $ do
     aliceUTxO <- genUTxOFor (genVerificationKey `genForParty` Fixture.alice)
     bobUTxO <- genUTxOFor (genVerificationKey `genForParty` Fixture.bob)
     carolUTxO <- genUTxOFor (genVerificationKey `genForParty` Fixture.carol)
@@ -419,46 +381,48 @@ snapshotUTxO snapshot =
 
 -- | A correctly signed snapshot. Given a snapshot number a snapshot signed by
 -- all participants (alice, bob and carol) with some UTxO contained is produced.
-signedSnapshot :: ModelSnapshot -> (Snapshot Tx, MultiSignature (Snapshot Tx))
-signedSnapshot ms =
+signedSnapshot :: SnapshotNumber -> (Snapshot Tx, MultiSignature (Snapshot Tx))
+signedSnapshot number =
   (snapshot, signatures)
  where
   snapshot =
     Snapshot
       { headId = mkHeadId Fixture.testPolicyId
-      , number = snapshotNumber ms
+      , number
       , confirmed = []
       , utxo = allUTxO
-      , utxoToDecommit = decommitUTxO
+      , utxoToDecommit = Nothing -- decommitUTxO
       }
-  (allUTxO, decommitUTxO) = pickUTxOToDecommit $ snapshotUTxO ms
+  allUTxO = snapshotUTxO number
+  -- (allUTxO, decommitUTxO) = pickUTxOToDecommit $ snapshotUTxO ms
 
   signatures = aggregate [sign sk snapshot | sk <- [Fixture.aliceSk, Fixture.bobSk, Fixture.carolSk]]
 
 splitHeadUTxO :: UTxO -> (UTxO, UTxO)
 splitHeadUTxO allUTxO =
-  let (headIn, headOut) = List.head $ List.filter (isHeadOutput . snd) (UTxO.pairs allUTxO)
-   in (UTxO.singleton (headIn, headOut), UTxO.filter (/= headOut) allUTxO)
+  let headUTxO = UTxO.filter isHeadOutput allUTxO
+   in (headUTxO, UTxO.filter (\o -> o `notElem` (snd <$> UTxO.pairs headUTxO)) allUTxO)
 
 pickUTxOToDecommit :: UTxO -> (UTxO, Maybe UTxO)
 pickUTxOToDecommit utxo = do
-  let pairs = UTxO.pairs utxo
+  let (headUTxO, rest) = splitHeadUTxO utxo
+  let pairs = UTxO.pairs rest
   case pairs of
     [] -> (utxo, Nothing)
     _ -> do
       let toDecommit = elements pairs `generateWith` 42
-      (UTxO.fromPairs $ filter (/= toDecommit) pairs, Just $ UTxO.singleton toDecommit)
+      (headUTxO <> UTxO.fromPairs (filter (/= toDecommit) pairs), Just $ UTxO.singleton toDecommit)
 
 -- | A confirmed snapshot (either initial or later confirmed), based on
 -- 'signedSnapshot'.
-confirmedSnapshot :: ModelSnapshot -> ConfirmedSnapshot Tx
+confirmedSnapshot :: SnapshotNumber -> ConfirmedSnapshot Tx
 confirmedSnapshot = \case
   0 ->
     InitialSnapshot
       { -- -- NOTE: The close validator would not check headId on close with
         -- initial snapshot, but we need to provide it still.
         headId = mkHeadId Fixture.testPolicyId
-      , initialUTxO = snapshotUTxO (Normal 0)
+      , initialUTxO = snapshotUTxO 0
       }
   number -> ConfirmedSnapshot{snapshot, signatures}
    where
@@ -480,7 +444,7 @@ openHeadUTxO =
     mkTxOutDatumInline
       Head.Open
         { parties = partyToChain <$> [Fixture.alice, Fixture.bob, Fixture.carol]
-        , utxoHash = toBuiltin $ hashUTxO @Tx $ snapshotUTxO (Normal 0)
+        , utxoHash = toBuiltin $ hashUTxO @Tx $ snapshotUTxO 0
         , contestationPeriod = CP.toChain Fixture.cperiod
         , headId = headIdToCurrencySymbol $ mkHeadId Fixture.testPolicyId
         , snapshotNumber = 0
@@ -540,7 +504,7 @@ newContestTx actor snapshot = do
 -- | Creates a fanout transaction using given utxo. NOTE: This uses fixtures for
 -- seedTxIn and contestation period. Consequently, the lower bound used is
 -- precisely at the maximum deadline slot as if everyone contested.
-newFanoutTx :: Actor -> ModelSnapshot -> AppM (Either FanoutTxError Tx)
+newFanoutTx :: Actor -> SnapshotNumber -> AppM (Either FanoutTxError Tx)
 newFanoutTx actor snapshot = do
   spendableUTxO <- get
   let (snapshot', _) = signedSnapshot snapshot