diff --git a/breacharbiter.go b/breacharbiter.go
index 6584cc17b..3bb334537 100644
--- a/breacharbiter.go
+++ b/breacharbiter.go
@@ -821,7 +821,7 @@ func (b *breachArbiter) handleBreachHandoff(breachEvent *ContractBreachEvent) {
 type breachedOutput struct {
 	amt         btcutil.Amount
 	outpoint    wire.OutPoint
-	witnessType input.WitnessType
+	witnessType input.StandardWitnessType
 	signDesc    input.SignDescriptor
 	confHeight  uint32
 
@@ -833,7 +833,7 @@ type breachedOutput struct {
 // makeBreachedOutput assembles a new breachedOutput that can be used by the
 // breach arbiter to construct a justice or sweep transaction.
 func makeBreachedOutput(outpoint *wire.OutPoint,
-	witnessType input.WitnessType,
+	witnessType input.StandardWitnessType,
 	secondLevelScript []byte,
 	signDescriptor *input.SignDescriptor,
 	confHeight uint32) breachedOutput {
@@ -883,9 +883,7 @@ func (bo *breachedOutput) CraftInputScript(signer input.Signer, txn *wire.MsgTx,
 
 	// First, we ensure that the witness generation function has been
 	// initialized for this breached output.
-	bo.witnessFunc = bo.witnessType.GenWitnessFunc(
-		signer, bo.SignDesc(),
-	)
+	bo.witnessFunc = bo.witnessType.WitnessGenerator(signer, bo.SignDesc())
 
 	// Now that we have ensured that the witness generation function has
 	// been initialized, we can proceed to execute it and generate the
@@ -993,7 +991,7 @@ func newRetributionInfo(chanPoint *wire.OutPoint,
 		// Using the breachedHtlc's incoming flag, determine the
 		// appropriate witness type that needs to be generated in order
 		// to sweep the HTLC output.
-		var htlcWitnessType input.WitnessType
+		var htlcWitnessType input.StandardWitnessType
 		if breachedHtlc.IsIncoming {
 			htlcWitnessType = input.HtlcAcceptedRevoke
 		} else {
@@ -1051,32 +1049,15 @@ func (b *breachArbiter) createJusticeTx(
 		// Grab locally scoped reference to breached output.
 		inp := &r.breachedOutputs[i]
 
-		// First, select the appropriate estimated witness weight for
+		// First, determine the appropriate estimated witness weight for
 		// the give witness type of this breached output. If the witness
-		// type is unrecognized, we will omit it from the transaction.
-		var witnessWeight int
-		switch inp.WitnessType() {
-		case input.CommitSpendNoDelayTweakless:
-			fallthrough
-		case input.CommitmentNoDelay:
-			witnessWeight = input.P2WKHWitnessSize
-
-		case input.CommitmentRevoke:
-			witnessWeight = input.ToLocalPenaltyWitnessSize
-
-		case input.HtlcOfferedRevoke:
-			witnessWeight = input.OfferedHtlcPenaltyWitnessSize
-
-		case input.HtlcAcceptedRevoke:
-			witnessWeight = input.AcceptedHtlcPenaltyWitnessSize
-
-		case input.HtlcSecondLevelRevoke:
-			witnessWeight = input.ToLocalPenaltyWitnessSize
-
-		default:
-			brarLog.Warnf("breached output in retribution info "+
-				"contains unexpected witness type: %v",
-				inp.WitnessType())
+		// weight cannot be estimated, we will omit it from the
+		// transaction.
+		witnessWeight, _, err := inp.WitnessType().SizeUpperBound()
+		if err != nil {
+			brarLog.Warnf("could not determine witness weight "+
+				"for breached output in retribution info: %v",
+				err)
 			continue
 		}
 		weightEstimate.AddWitnessInput(witnessWeight)
@@ -1555,7 +1536,7 @@ func (bo *breachedOutput) Decode(r io.Reader) error {
 	if _, err := io.ReadFull(r, scratch[:2]); err != nil {
 		return err
 	}
-	bo.witnessType = input.WitnessType(
+	bo.witnessType = input.StandardWitnessType(
 		binary.BigEndian.Uint16(scratch[:2]),
 	)
 
diff --git a/input/input.go b/input/input.go
index 10366b705..a63048cbf 100644
--- a/input/input.go
+++ b/input/input.go
@@ -114,9 +114,7 @@ func NewBaseInput(outpoint *wire.OutPoint, witnessType WitnessType,
 func (bi *BaseInput) CraftInputScript(signer Signer, txn *wire.MsgTx,
 	hashCache *txscript.TxSigHashes, txinIdx int) (*Script, error) {
 
-	witnessFunc := bi.witnessType.GenWitnessFunc(
-		signer, bi.SignDesc(),
-	)
+	witnessFunc := bi.witnessType.WitnessGenerator(signer, bi.SignDesc())
 
 	return witnessFunc(txn, hashCache, txinIdx)
 }
diff --git a/input/witnessgen.go b/input/witnessgen.go
index 28a02d086..1b8c930e1 100644
--- a/input/witnessgen.go
+++ b/input/witnessgen.go
@@ -7,87 +7,125 @@ import (
 	"github.com/btcsuite/btcd/wire"
 )
 
-// WitnessType determines how an output's witness will be generated. The
-// default commitmentTimeLock type will generate a witness that will allow
-// spending of a time-locked transaction enforced by CheckSequenceVerify.
-type WitnessType uint16
+// WitnessGenerator represents a function that is able to generate the final
+// witness for a particular public key script. Additionally, if required, this
+// function will also return the sigScript for spending nested P2SH witness
+// outputs. This function acts as an abstraction layer, hiding the details of
+// the underlying script.
+type WitnessGenerator func(tx *wire.MsgTx, hc *txscript.TxSigHashes,
+	inputIndex int) (*Script, error)
+
+// WitnessType determines how an output's witness will be generated. This
+// interface can be implemented to be used for custom sweep scripts if the
+// pre-defined StandardWitnessType list doesn't provide a suitable one.
+type WitnessType interface {
+	// String returns a human readable version of the WitnessType.
+	String() string
+
+	// WitnessGenerator will return a WitnessGenerator function that an
+	// output uses to generate the witness and optionally the sigScript for
+	// a sweep transaction.
+	WitnessGenerator(signer Signer,
+		descriptor *SignDescriptor) WitnessGenerator
+
+	// SizeUpperBound returns the maximum length of the witness of this
+	// WitnessType if it would be included in a tx. It also returns if the
+	// output itself is a nested p2sh output, if so then we need to take
+	// into account the extra sigScript data size.
+	SizeUpperBound() (int, bool, error)
+
+	// AddWeightEstimation adds the estimated size of the witness in bytes
+	// to the given weight estimator and returns the number of
+	// CSVs/CLTVs used by the script.
+	AddWeightEstimation(estimator *TxWeightEstimator) (int, int, error)
+}
+
+// StandardWitnessType is a numeric representation of standard pre-defined types
+// of witness configurations.
+type StandardWitnessType uint16
+
+// A compile time check to ensure StandardWitnessType implements the
+// WitnessType interface.
+var _ WitnessType = (StandardWitnessType)(0)
 
 const (
 	// CommitmentTimeLock is a witness that allows us to spend the output of
 	// a commitment transaction after a relative lock-time lockout.
-	CommitmentTimeLock WitnessType = 0
+	CommitmentTimeLock StandardWitnessType = 0
 
 	// CommitmentNoDelay is a witness that allows us to spend a settled
 	// no-delay output immediately on a counterparty's commitment
 	// transaction.
-	CommitmentNoDelay WitnessType = 1
+	CommitmentNoDelay StandardWitnessType = 1
 
 	// CommitmentRevoke is a witness that allows us to sweep the settled
 	// output of a malicious counterparty's who broadcasts a revoked
 	// commitment transaction.
-	CommitmentRevoke WitnessType = 2
+	CommitmentRevoke StandardWitnessType = 2
 
 	// HtlcOfferedRevoke is a witness that allows us to sweep an HTLC which
 	// we offered to the remote party in the case that they broadcast a
 	// revoked commitment state.
-	HtlcOfferedRevoke WitnessType = 3
+	HtlcOfferedRevoke StandardWitnessType = 3
 
 	// HtlcAcceptedRevoke is a witness that allows us to sweep an HTLC
 	// output sent to us in the case that the remote party broadcasts a
 	// revoked commitment state.
-	HtlcAcceptedRevoke WitnessType = 4
+	HtlcAcceptedRevoke StandardWitnessType = 4
 
 	// HtlcOfferedTimeoutSecondLevel is a witness that allows us to sweep
 	// an HTLC output that we extended to a party, but was never fulfilled.
 	// This HTLC output isn't directly on the commitment transaction, but
 	// is the result of a confirmed second-level HTLC transaction. As a
 	// result, we can only spend this after a CSV delay.
-	HtlcOfferedTimeoutSecondLevel WitnessType = 5
+	HtlcOfferedTimeoutSecondLevel StandardWitnessType = 5
 
 	// HtlcAcceptedSuccessSecondLevel is a witness that allows us to sweep
 	// an HTLC output that was offered to us, and for which we have a
 	// payment preimage. This HTLC output isn't directly on our commitment
 	// transaction, but is the result of confirmed second-level HTLC
 	// transaction. As a result, we can only spend this after a CSV delay.
-	HtlcAcceptedSuccessSecondLevel WitnessType = 6
+	HtlcAcceptedSuccessSecondLevel StandardWitnessType = 6
 
 	// HtlcOfferedRemoteTimeout is a witness that allows us to sweep an
 	// HTLC that we offered to the remote party which lies in the
 	// commitment transaction of the remote party. We can spend this output
 	// after the absolute CLTV timeout of the HTLC as passed.
-	HtlcOfferedRemoteTimeout WitnessType = 7
+	HtlcOfferedRemoteTimeout StandardWitnessType = 7
 
 	// HtlcAcceptedRemoteSuccess is a witness that allows us to sweep an
 	// HTLC that was offered to us by the remote party. We use this witness
 	// in the case that the remote party goes to chain, and we know the
 	// pre-image to the HTLC. We can sweep this without any additional
 	// timeout.
-	HtlcAcceptedRemoteSuccess WitnessType = 8
+	HtlcAcceptedRemoteSuccess StandardWitnessType = 8
 
 	// HtlcSecondLevelRevoke is a witness that allows us to sweep an HTLC
 	// from the remote party's commitment transaction in the case that the
 	// broadcast a revoked commitment, but then also immediately attempt to
 	// go to the second level to claim the HTLC.
-	HtlcSecondLevelRevoke WitnessType = 9
+	HtlcSecondLevelRevoke StandardWitnessType = 9
 
 	// WitnessKeyHash is a witness type that allows us to spend a regular
 	// p2wkh output that's sent to an output which is under complete
 	// control of the backing wallet.
-	WitnessKeyHash WitnessType = 10
+	WitnessKeyHash StandardWitnessType = 10
 
 	// NestedWitnessKeyHash is a witness type that allows us to sweep an
 	// output that sends to a nested P2SH script that pays to a key solely
 	// under our control. The witness generated needs to include the
-	NestedWitnessKeyHash WitnessType = 11
+	NestedWitnessKeyHash StandardWitnessType = 11
 
 	// CommitSpendNoDelayTweakless is similar to the CommitSpendNoDelay
 	// type, but it omits the tweak that randomizes the key we need to
 	// spend with a channel peer supplied set of randomness.
-	CommitSpendNoDelayTweakless = 12
+	CommitSpendNoDelayTweakless StandardWitnessType = 12
 )
 
-// Stirng returns a human readable version of the target WitnessType.
-func (wt WitnessType) String() string {
+// String returns a human readable version of the target WitnessType.
+//
+// NOTE: This is part of the WitnessType interface.
+func (wt StandardWitnessType) String() string {
 	switch wt {
 	case CommitmentTimeLock:
 		return "CommitmentTimeLock"
@@ -127,19 +165,13 @@ func (wt WitnessType) String() string {
 	}
 }
 
-// WitnessGenerator represents a function which is able to generate the final
-// witness for a particular public key script. Additionally, if required, this
-// function will also return the sigScript for spending nested P2SH witness
-// outputs. This function acts as an abstraction layer, hiding the details of
-// the underlying script.
-type WitnessGenerator func(tx *wire.MsgTx, hc *txscript.TxSigHashes,
-	inputIndex int) (*Script, error)
-
-// GenWitnessFunc will return a WitnessGenerator function that an output uses
+// WitnessGenerator will return a WitnessGenerator function that an output uses
 // to generate the witness and optionally the sigScript for a sweep
 // transaction. The sigScript will be generated if the witness type warrants
 // one for spending, such as the NestedWitnessKeyHash witness type.
-func (wt WitnessType) GenWitnessFunc(signer Signer,
+//
+// NOTE: This is part of the WitnessType interface.
+func (wt StandardWitnessType) WitnessGenerator(signer Signer,
 	descriptor *SignDescriptor) WitnessGenerator {
 
 	return func(tx *wire.MsgTx, hc *txscript.TxSigHashes,
@@ -262,5 +294,116 @@ func (wt WitnessType) GenWitnessFunc(signer Signer,
 			return nil, fmt.Errorf("unknown witness type: %v", wt)
 		}
 	}
-
+}
+
+// SizeUpperBound returns the maximum length of the witness of this witness
+// type if it would be included in a tx. We also return if the output itself is
+// a nested p2sh output, if so then we need to take into account the extra
+// sigScript data size.
+//
+// NOTE: This is part of the WitnessType interface.
+func (wt StandardWitnessType) SizeUpperBound() (int, bool, error) {
+	switch wt {
+
+	// Outputs on a remote commitment transaction that pay directly to us.
+	case CommitSpendNoDelayTweakless:
+		fallthrough
+	case WitnessKeyHash:
+		fallthrough
+	case CommitmentNoDelay:
+		return P2WKHWitnessSize, false, nil
+
+	// Outputs on a past commitment transaction that pay directly
+	// to us.
+	case CommitmentTimeLock:
+		return ToLocalTimeoutWitnessSize, false, nil
+
+	// Outgoing second layer HTLC's that have confirmed within the
+	// chain, and the output they produced is now mature enough to
+	// sweep.
+	case HtlcOfferedTimeoutSecondLevel:
+		return ToLocalTimeoutWitnessSize, false, nil
+
+	// Incoming second layer HTLC's that have confirmed within the
+	// chain, and the output they produced is now mature enough to
+	// sweep.
+	case HtlcAcceptedSuccessSecondLevel:
+		return ToLocalTimeoutWitnessSize, false, nil
+
+	// An HTLC on the commitment transaction of the remote party,
+	// that has had its absolute timelock expire.
+	case HtlcOfferedRemoteTimeout:
+		return AcceptedHtlcTimeoutWitnessSize, false, nil
+
+	// An HTLC on the commitment transaction of the remote party,
+	// that can be swept with the preimage.
+	case HtlcAcceptedRemoteSuccess:
+		return OfferedHtlcSuccessWitnessSize, false, nil
+
+	// A nested P2SH input that has a p2wkh witness script. We'll mark this
+	// as nested P2SH so the caller can estimate the weight properly
+	// including the sigScript.
+	case NestedWitnessKeyHash:
+		return P2WKHWitnessSize, true, nil
+
+	// The revocation output on a revoked commitment transaction.
+	case CommitmentRevoke:
+		return ToLocalPenaltyWitnessSize, false, nil
+
+	// The revocation output on a revoked HTLC that we offered to the remote
+	// party.
+	case HtlcOfferedRevoke:
+		return OfferedHtlcPenaltyWitnessSize, false, nil
+
+	// The revocation output on a revoked HTLC that was sent to us.
+	case HtlcAcceptedRevoke:
+		return AcceptedHtlcPenaltyWitnessSize, false, nil
+
+	// The revocation output of a second level output of an HTLC.
+	case HtlcSecondLevelRevoke:
+		return ToLocalPenaltyWitnessSize, false, nil
+	}
+
+	return 0, false, fmt.Errorf("unexpected witness type: %v", wt)
+}
+
+// AddWeightEstimation adds the estimated size of the witness in bytes to the
+// given weight estimator and returns the number of CSVs/CLTVs used by the
+// script.
+//
+// NOTE: This is part of the WitnessType interface.
+func (wt StandardWitnessType) AddWeightEstimation(
+	estimator *TxWeightEstimator) (int, int, error) {
+
+	var (
+		csvCount  = 0
+		cltvCount = 0
+	)
+
+	// For fee estimation purposes, we'll now attempt to obtain an
+	// upper bound on the weight this input will add when fully
+	// populated.
+	size, isNestedP2SH, err := wt.SizeUpperBound()
+	if err != nil {
+		return 0, 0, err
+	}
+
+	// If this is a nested P2SH input, then we'll need to factor in
+	// the additional data push within the sigScript.
+	if isNestedP2SH {
+		estimator.AddNestedP2WSHInput(size)
+	} else {
+		estimator.AddWitnessInput(size)
+	}
+
+	switch wt {
+	case CommitmentTimeLock,
+		HtlcOfferedTimeoutSecondLevel,
+		HtlcAcceptedSuccessSecondLevel:
+		csvCount++
+	case HtlcOfferedRemoteTimeout:
+		cltvCount++
+	}
+
+	return csvCount, cltvCount, nil
 }
diff --git a/server.go b/server.go
index 6297fc855..5aa496622 100644
--- a/server.go
+++ b/server.go
@@ -764,7 +764,7 @@ func newServer(listenAddrs []net.Addr, chanDB *channeldb.DB,
 		sweep.DefaultBatchWindowDuration)
 
 	sweeperStore, err := sweep.NewSweeperStore(
-		chanDB, activeNetParams.GenesisHash,
+		chanDB.DB, activeNetParams.GenesisHash,
 	)
 	if err != nil {
 		srvrLog.Errorf("unable to create sweeper store: %v", err)
@@ -780,7 +780,6 @@ func newServer(listenAddrs []net.Addr, chanDB *channeldb.DB,
 			return time.NewTimer(sweep.DefaultBatchWindowDuration).C
 		},
 		Notifier:             cc.chainNotifier,
-		ChainIO:              cc.chainIO,
 		Store:                sweeperStore,
 		MaxInputsPerTx:       sweep.DefaultMaxInputsPerTx,
 		MaxSweepAttempts:     sweep.DefaultMaxSweepAttempts,
diff --git a/sweep/store.go b/sweep/store.go
index ef6ba99c9..f1ca84759 100644
--- a/sweep/store.go
+++ b/sweep/store.go
@@ -9,7 +9,6 @@ import (
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/coreos/bbolt"
-	"github.com/lightningnetwork/lnd/channeldb"
 )
 
 var (
@@ -57,11 +56,11 @@ type SweeperStore interface {
 }
 
 type sweeperStore struct {
-	db *channeldb.DB
+	db *bbolt.DB
 }
 
 // NewSweeperStore returns a new store instance.
-func NewSweeperStore(db *channeldb.DB, chainHash *chainhash.Hash) (
+func NewSweeperStore(db *bbolt.DB, chainHash *chainhash.Hash) (
 	SweeperStore, error) {
 
 	err := db.Update(func(tx *bbolt.Tx) error {
diff --git a/sweep/store_test.go b/sweep/store_test.go
index 23714c781..6853c7bb0 100644
--- a/sweep/store_test.go
+++ b/sweep/store_test.go
@@ -53,7 +53,7 @@ func TestStore(t *testing.T) {
 
 		testStore(t, func() (SweeperStore, error) {
 			var chain chainhash.Hash
-			return NewSweeperStore(cdb, &chain)
+			return NewSweeperStore(cdb.DB, &chain)
 		})
 	})
 	t.Run("mock", func(t *testing.T) {
diff --git a/sweep/sweeper.go b/sweep/sweeper.go
index 03fcf6ca7..2e3a1e1ca 100644
--- a/sweep/sweeper.go
+++ b/sweep/sweeper.go
@@ -212,9 +212,6 @@ type UtxoSweeperConfig struct {
 	// certain on-chain events.
 	Notifier chainntnfs.ChainNotifier
 
-	// ChainIO is used  to determine the current block height.
-	ChainIO lnwallet.BlockChainIO
-
 	// Store stores the published sweeper txes.
 	Store SweeperStore
 
@@ -323,20 +320,10 @@ func (s *UtxoSweeper) Start() error {
 	// not change from here on.
 	s.relayFeeRate = s.cfg.FeeEstimator.RelayFeePerKW()
 
-	// Register for block epochs to retry sweeping every block.
-	bestHash, bestHeight, err := s.cfg.ChainIO.GetBestBlock()
-	if err != nil {
-		return fmt.Errorf("get best block: %v", err)
-	}
-
-	log.Debugf("Best height: %v", bestHeight)
-
-	blockEpochs, err := s.cfg.Notifier.RegisterBlockEpochNtfn(
-		&chainntnfs.BlockEpoch{
-			Height: bestHeight,
-			Hash:   bestHash,
-		},
-	)
+	// We need to register for block epochs and retry sweeping every block.
+	// We should get a notification with the current best block immediately
+	// if we don't provide any epoch. We'll wait for that in the collector.
+	blockEpochs, err := s.cfg.Notifier.RegisterBlockEpochNtfn(nil)
 	if err != nil {
 		return fmt.Errorf("register block epoch ntfn: %v", err)
 	}
@@ -347,10 +334,7 @@ func (s *UtxoSweeper) Start() error {
 		defer blockEpochs.Cancel()
 		defer s.wg.Done()
 
-		err := s.collector(blockEpochs.Epochs, bestHeight)
-		if err != nil {
-			log.Errorf("sweeper stopped: %v", err)
-		}
+		s.collector(blockEpochs.Epochs)
 	}()
 
 	return nil
@@ -445,8 +429,18 @@ func (s *UtxoSweeper) feeRateForPreference(
 
 // collector is the sweeper main loop. It processes new inputs, spend
 // notifications and counts down to publication of the sweep tx.
-func (s *UtxoSweeper) collector(blockEpochs <-chan *chainntnfs.BlockEpoch,
-	bestHeight int32) error {
+func (s *UtxoSweeper) collector(blockEpochs <-chan *chainntnfs.BlockEpoch) {
+	// We registered for the block epochs with a nil request. The notifier
+	// should send us the current best block immediately. So we need to wait
+	// for it here because we need to know the current best height.
+	var bestHeight int32
+	select {
+	case bestBlock := <-blockEpochs:
+		bestHeight = bestBlock.Height
+
+	case <-s.quit:
+		return
+	}
 
 	for {
 		select {
@@ -622,7 +616,7 @@ func (s *UtxoSweeper) collector(blockEpochs <-chan *chainntnfs.BlockEpoch,
 		// sweep.
 		case epoch, ok := <-blockEpochs:
 			if !ok {
-				return nil
+				return
 			}
 
 			bestHeight = epoch.Height
@@ -635,7 +629,7 @@ func (s *UtxoSweeper) collector(blockEpochs <-chan *chainntnfs.BlockEpoch,
 			}
 
 		case <-s.quit:
-			return nil
+			return
 		}
 	}
 }
diff --git a/sweep/sweeper_test.go b/sweep/sweeper_test.go
index b4139fbac..501ec3f87 100644
--- a/sweep/sweeper_test.go
+++ b/sweep/sweeper_test.go
@@ -130,9 +130,8 @@ func createSweeperTestContext(t *testing.T) *sweeperTestContext {
 			ctx.timeoutChan <- c
 			return c
 		},
-		Store:   store,
-		Signer:  &mockSigner{},
-		ChainIO: &mockChainIO{},
+		Store:  store,
+		Signer: &mockSigner{},
 		GenSweepScript: func() ([]byte, error) {
 			script := []byte{outputScriptCount}
 			outputScriptCount++
diff --git a/sweep/test_utils.go b/sweep/test_utils.go
index 46ee6dc3d..df44cd105 100644
--- a/sweep/test_utils.go
+++ b/sweep/test_utils.go
@@ -10,12 +10,12 @@ import (
 	"github.com/btcsuite/btcd/wire"
 	"github.com/lightningnetwork/lnd/chainntnfs"
 	"github.com/lightningnetwork/lnd/input"
-	"github.com/lightningnetwork/lnd/lnwallet"
 )
 
 var (
 	defaultTestTimeout = 5 * time.Second
-	mockChainIOHeight  = int32(100)
+	mockChainHash, _   = chainhash.NewHashFromStr("00aabbccddeeff")
+	mockChainHeight    = int32(100)
 )
 
 type mockSigner struct {
@@ -155,12 +155,22 @@ func (m *MockNotifier) RegisterBlockEpochNtfn(
 	log.Tracef("Mock block ntfn registered")
 
 	m.mutex.Lock()
-	epochChan := make(chan *chainntnfs.BlockEpoch, 0)
-	bestHeight := int32(0)
-	if bestBlock != nil {
-		bestHeight = bestBlock.Height
+	epochChan := make(chan *chainntnfs.BlockEpoch, 1)
+
+	// The real notifier returns a notification with the current block hash
+	// and height immediately if no best block hash or height is specified
+	// in the request. We want to emulate this behaviour as well for the
+	// mock.
+	switch {
+	case bestBlock == nil:
+		epochChan <- &chainntnfs.BlockEpoch{
+			Hash:   mockChainHash,
+			Height: mockChainHeight,
+		}
+		m.epochChan[epochChan] = mockChainHeight
+	default:
+		m.epochChan[epochChan] = bestBlock.Height
 	}
-	m.epochChan[epochChan] = bestHeight
 	m.mutex.Unlock()
 
 	return &chainntnfs.BlockEpochEvent{
@@ -235,25 +245,3 @@ func (m *MockNotifier) RegisterSpendNtfn(outpoint *wire.OutPoint,
 		},
 	}, nil
 }
-
-type mockChainIO struct{}
-
-var _ lnwallet.BlockChainIO = (*mockChainIO)(nil)
-
-func (m *mockChainIO) GetBestBlock() (*chainhash.Hash, int32, error) {
-	return nil, mockChainIOHeight, nil
-}
-
-func (m *mockChainIO) GetUtxo(op *wire.OutPoint, pkScript []byte,
-	heightHint uint32, _ <-chan struct{}) (*wire.TxOut, error) {
-
-	return nil, nil
-}
-
-func (m *mockChainIO) GetBlockHash(blockHeight int64) (*chainhash.Hash, error) {
-	return nil, nil
-}
-
-func (m *mockChainIO) GetBlock(blockHash *chainhash.Hash) (*wire.MsgBlock, error) {
-	return nil, nil
-}
diff --git a/sweep/txgenerator.go b/sweep/txgenerator.go
index a06dad963..4e383c701 100644
--- a/sweep/txgenerator.go
+++ b/sweep/txgenerator.go
@@ -56,7 +56,7 @@ func generateInputPartitionings(sweepableInputs []input.Input,
 	// on the signature length, which is not known yet at this point.
 	yields := make(map[wire.OutPoint]int64)
 	for _, input := range sweepableInputs {
-		size, _, err := getInputWitnessSizeUpperBound(input)
+		size, _, err := input.WitnessType().SizeUpperBound()
 		if err != nil {
 			return nil, fmt.Errorf(
 				"failed adding input weight: %v", err)
@@ -126,7 +126,7 @@ func getPositiveYieldInputs(sweepableInputs []input.Input, maxInputs int,
 	for idx, input := range sweepableInputs {
 		// Can ignore error, because it has already been checked when
 		// calculating the yields.
-		size, isNestedP2SH, _ := getInputWitnessSizeUpperBound(input)
+		size, isNestedP2SH, _ := input.WitnessType().SizeUpperBound()
 
 		// Keep a running weight estimate of the input set.
 		if isNestedP2SH {
@@ -251,59 +251,6 @@ func createSweepTx(inputs []input.Input, outputPkScript []byte,
 	return sweepTx, nil
 }
 
-// getInputWitnessSizeUpperBound returns the maximum length of the witness for
-// the given input if it would be included in a tx. We also return if the
-// output itself is a nested p2sh output, if so then we need to take into
-// account the extra sigScript data size.
-func getInputWitnessSizeUpperBound(inp input.Input) (int, bool, error) {
-	switch inp.WitnessType() {
-
-	// Outputs on a remote commitment transaction that pay directly to us.
-	case input.CommitSpendNoDelayTweakless:
-		fallthrough
-	case input.WitnessKeyHash:
-		fallthrough
-	case input.CommitmentNoDelay:
-		return input.P2WKHWitnessSize, false, nil
-
-	// Outputs on a past commitment transaction that pay directly
-	// to us.
-	case input.CommitmentTimeLock:
-		return input.ToLocalTimeoutWitnessSize, false, nil
-
-	// Outgoing second layer HTLC's that have confirmed within the
-	// chain, and the output they produced is now mature enough to
-	// sweep.
-	case input.HtlcOfferedTimeoutSecondLevel:
-		return input.ToLocalTimeoutWitnessSize, false, nil
-
-	// Incoming second layer HTLC's that have confirmed within the
-	// chain, and the output they produced is now mature enough to
-	// sweep.
-	case input.HtlcAcceptedSuccessSecondLevel:
-		return input.ToLocalTimeoutWitnessSize, false, nil
-
-	// An HTLC on the commitment transaction of the remote party,
-	// that has had its absolute timelock expire.
-	case input.HtlcOfferedRemoteTimeout:
-		return input.AcceptedHtlcTimeoutWitnessSize, false, nil
-
-	// An HTLC on the commitment transaction of the remote party,
-	// that can be swept with the preimage.
-	case input.HtlcAcceptedRemoteSuccess:
-		return input.OfferedHtlcSuccessWitnessSize, false, nil
-
-	// A nested P2SH input that has a p2wkh witness script. We'll mark this
-	// as nested P2SH so the caller can estimate the weight properly
-	// including the sigScript.
-	case input.NestedWitnessKeyHash:
-		return input.P2WKHWitnessSize, true, nil
-	}
-
-	return 0, false, fmt.Errorf("unexpected witness type: %v",
-		inp.WitnessType())
-}
-
 // getWeightEstimate returns a weight estimate for the given inputs.
 // Additionally, it returns counts for the number of csv and cltv inputs.
 func getWeightEstimate(inputs []input.Input) ([]input.Input, int64, int, int) {
@@ -328,10 +275,8 @@ func getWeightEstimate(inputs []input.Input) ([]input.Input, int64, int, int) {
 	for i := range inputs {
 		inp := inputs[i]
 
-		// For fee estimation purposes, we'll now attempt to obtain an
-		// upper bound on the weight this input will add when fully
-		// populated.
-		size, isNestedP2SH, err := getInputWitnessSizeUpperBound(inp)
+		wt := inp.WitnessType()
+		inpCsv, inpCltv, err := wt.AddWeightEstimation(&weightEstimate)
 		if err != nil {
 			log.Warn(err)
 
@@ -339,23 +284,8 @@ func getWeightEstimate(inputs []input.Input) ([]input.Input, int64, int, int) {
 			// given.
 			continue
 		}
-
-		// If this is a nested P2SH input, then we'll need to factor in
-		// the additional data push within the sigScript.
-		if isNestedP2SH {
-			weightEstimate.AddNestedP2WSHInput(size)
-		} else {
-			weightEstimate.AddWitnessInput(size)
-		}
-
-		switch inp.WitnessType() {
-		case input.CommitmentTimeLock,
-			input.HtlcOfferedTimeoutSecondLevel,
-			input.HtlcAcceptedSuccessSecondLevel:
-			csvCount++
-		case input.HtlcOfferedRemoteTimeout:
-			cltvCount++
-		}
+		csvCount += inpCsv
+		cltvCount += inpCltv
 		sweepInputs = append(sweepInputs, inp)
 	}
 
diff --git a/sweep/txgenerator_test.go b/sweep/txgenerator_test.go
new file mode 100644
index 000000000..571eb134f
--- /dev/null
+++ b/sweep/txgenerator_test.go
@@ -0,0 +1,49 @@
+package sweep
+
+import (
+	"testing"
+
+	"github.com/btcsuite/btcd/wire"
+	"github.com/lightningnetwork/lnd/input"
+)
+
+var (
+	witnessTypes = []input.WitnessType{
+		input.CommitmentTimeLock,
+		input.HtlcAcceptedSuccessSecondLevel,
+		input.HtlcOfferedRemoteTimeout,
+		input.WitnessKeyHash,
+	}
+	expectedWeight = int64(1459)
+	expectedCsv    = 2
+	expectedCltv   = 1
+)
+
+// TestWeightEstimate tests that the estimated weight and number of CSVs/CLTVs
+// used is correct for a transaction that uses inputs with the witness types
+// defined in witnessTypes.
+func TestWeightEstimate(t *testing.T) {
+	t.Parallel()
+
+	var inputs []input.Input
+	for _, witnessType := range witnessTypes {
+		inputs = append(inputs, input.NewBaseInput(
+			&wire.OutPoint{}, witnessType,
+			&input.SignDescriptor{}, 0,
+		))
+	}
+
+	_, weight, csv, cltv := getWeightEstimate(inputs)
+	if weight != expectedWeight {
+		t.Fatalf("unexpected weight. expected %d but got %d.",
+			expectedWeight, weight)
+	}
+	if csv != expectedCsv {
+		t.Fatalf("unexpected csv count. expected %d but got %d.",
+			expectedCsv, csv)
+	}
+	if cltv != expectedCltv {
+		t.Fatalf("unexpected cltv count. expected %d but got %d.",
+			expectedCltv, cltv)
+	}
+}
diff --git a/utxonursery.go b/utxonursery.go
index 713c8c9d4..b08cf0e2a 100644
--- a/utxonursery.go
+++ b/utxonursery.go
@@ -1344,7 +1344,7 @@ type kidOutput struct {
 }
 
 func makeKidOutput(outpoint, originChanPoint *wire.OutPoint,
-	blocksToMaturity uint32, witnessType input.WitnessType,
+	blocksToMaturity uint32, witnessType input.StandardWitnessType,
 	signDescriptor *input.SignDescriptor,
 	absoluteMaturity uint32) kidOutput {
 
@@ -1423,7 +1423,7 @@ func (k *kidOutput) Encode(w io.Writer) error {
 		return err
 	}
 
-	byteOrder.PutUint16(scratch[:2], uint16(k.WitnessType()))
+	byteOrder.PutUint16(scratch[:2], uint16(k.witnessType))
 	if _, err := w.Write(scratch[:2]); err != nil {
 		return err
 	}
@@ -1473,7 +1473,7 @@ func (k *kidOutput) Decode(r io.Reader) error {
 	if _, err := r.Read(scratch[:2]); err != nil {
 		return err
 	}
-	k.witnessType = input.WitnessType(byteOrder.Uint16(scratch[:2]))
+	k.witnessType = input.StandardWitnessType(byteOrder.Uint16(scratch[:2]))
 
 	return input.ReadSignDescriptor(r, &k.signDesc)
 }