func (s *txnSuite) TestRetryHooks(c *gc.C) {
	s.insertDoc(c, "1", "Foo")
	defer txntesting.SetRetryHooks(c, s.txnRunner, func() {
		s.setDocName(c, "1", "Bar")
	}, func() {
		s.setDocName(c, "1", "Foo")
	}).Check()

	maxAttempt := 0
	buildTxn := func(attempt int) ([]txn.Op, error) {
		maxAttempt = attempt
		ops := []txn.Op{{
			C:      s.collection.Name,
			Id:     "1",
			Assert: bson.D{{"name", "Foo"}},
			Update: bson.D{{"$set", bson.D{{"name", "FooBar"}}}},
		}}
		return ops, nil
	}
	err := s.txnRunner.Run(buildTxn)
	c.Assert(err, gc.IsNil)
	c.Assert(maxAttempt, gc.Equals, 2)
	var found simpleDoc
	err = s.collection.FindId("1").One(&found)
	c.Assert(err, gc.IsNil)
	doc := simpleDoc{"1", "FooBar"}
	c.Assert(found, gc.DeepEquals, doc)
}

func (s *ClientTrickyRaceSuite) TestExtendLease_WorksDespite_ExpireThenReclaimSameHolder(c *gc.C) {

	// Set up hooks to expire and reclaim the lease before the extend gets a
	// chance; and to verify that the second attempt successfully extends.
	defer txntesting.SetRetryHooks(c, s.sut.Runner, func() {
		s.blocker.Clock.Advance(90 * time.Second)
		err := s.blocker.Client.ExpireLease("name")
		c.Check(err, jc.ErrorIsNil)
		err = s.blocker.Client.ClaimLease("name", lease.Request{"holder", time.Minute})
		c.Check(err, jc.ErrorIsNil)
	}, func() {
		err := s.blocker.Client.Refresh()
		c.Check(err, jc.ErrorIsNil)
		c.Check("name", s.blocker.Expiry(), s.blocker.Zero.Add(5*time.Minute))
	})()

	// Try to extend; check it worked.
	err := s.sut.Client.ExtendLease("name", lease.Request{"holder", 5 * time.Minute})
	c.Check(err, jc.ErrorIsNil)
}

func (s *ClientTrickyRaceSuite) TestExtendLease_WorksDespite_ShorterExtendLease(c *gc.C) {

	shorterRequest := 90 * time.Second
	longerRequest := 120 * time.Second

	// Set up hooks to extend the lease by a little, before the SUT's extend
	// gets a chance; and then to verify state after it's applied its retry.
	defer txntesting.SetRetryHooks(c, s.sut.Runner, func() {
		err := s.blocker.Client.ExtendLease("name", lease.Request{"holder", shorterRequest})
		c.Check(err, jc.ErrorIsNil)
	}, func() {
		err := s.blocker.Client.Refresh()
		c.Check(err, jc.ErrorIsNil)
		c.Check("name", s.blocker.Expiry(), s.blocker.Zero.Add(longerRequest))
	})()

	// Extend the lease.
	err := s.sut.Client.ExtendLease("name", lease.Request{"holder", longerRequest})
	c.Check(err, jc.ErrorIsNil)
}

func SetRetryHooks(c *gc.C, st *State, block, check func()) txntesting.TransactionChecker {
	return txntesting.SetRetryHooks(c, newRunnerForHooks(st), block, check)
}

func SetRetryHooks(c *gc.C, st *State, block, check func()) txntesting.TransactionChecker {
	runner := jujutxn.NewRunner(jujutxn.RunnerParams{Database: st.db})
	st.transactionRunner = runner
	return txntesting.SetRetryHooks(c, runner, block, check)
}

func SetRetryHooks(c *gc.C, st *State, block, check func()) txntesting.TransactionChecker {
	return txntesting.SetRetryHooks(c, st.transactionRunner, block, check)
}