func (s *txnSuite) TestRunTransactionObserver(c *gc.C) {
	type args struct {
		ops []txn.Op
		err error
	}
	var calls []args
	runner := jujutxn.NewRunner(jujutxn.RunnerParams{
		RunTransactionObserver: func(ops []txn.Op, err error) {
			calls = append(calls, args{ops, err})
		},
	})
	fake := &fakeRunner{errors: []error{
		txn.ErrAborted,
		nil,
	}}
	jujutxn.SetRunnerFunc(runner, fake.new)
	ops := []txn.Op{{}}
	buildTxn := func(attempt int) ([]txn.Op, error) {
		return ops, nil
	}
	err := runner.Run(buildTxn)
	c.Check(err, gc.IsNil)
	c.Check(calls, gc.HasLen, 2)
	c.Check(calls[0].ops, gc.DeepEquals, ops)
	c.Check(calls[0].err, gc.Equals, txn.ErrAborted)
	c.Check(calls[1].ops, gc.DeepEquals, ops)
	c.Check(calls[1].err, gc.IsNil)
}

func newMultiEnvRunner(envUUID string, db *mgo.Database, assertEnvAlive bool) jujutxn.Runner {
	return &multiEnvRunner{
		rawRunner:      jujutxn.NewRunner(jujutxn.RunnerParams{Database: db}),
		envUUID:        envUUID,
		assertEnvAlive: assertEnvAlive,
	}
}

func (s *txnSuite) SetUpTest(c *gc.C) {
	s.IsolationSuite.SetUpTest(c)
	s.MgoSuite.SetUpTest(c)
	db := s.Session.DB("juju")
	s.collection = db.C("test")
	s.txnRunner = jujutxn.NewRunner(jujutxn.RunnerParams{Database: db})
}

// rawTxnRunner returns a transaction runner that won't perform
// automatic addition of environment UUIDs into transaction
// operations, even for collections that contain documents for
// multiple environments. It should be used rarely.
func (st *State) rawTxnRunner(session *mgo.Session) jujutxn.Runner {
	if st.transactionRunner != nil {
		return getRawRunner(st.transactionRunner)
	}
	return jujutxn.NewRunner(jujutxn.RunnerParams{
		Database: st.db.With(session),
	})
}

func NewTestMongo(database *mgo.Database) *TestMongo {
	return &TestMongo{
		database: database,
		runner: txn.NewRunner(txn.RunnerParams{
			Database: database,
		}),
	}
}

// NewMongo returns a *Mongo backed by the supplied database.
func NewMongo(database *mgo.Database) *Mongo {
	return &Mongo{
		database: database,
		runner: jujutxn.NewRunner(jujutxn.RunnerParams{
			Database: database,
		}),
	}
}

func newStorage(st *State, uuid, metadataCollection string) binarystorage.StorageCloser {
	session := st.session.Clone()
	rs := blobstore.NewGridFS(blobstoreDB, uuid, session)
	db := session.DB(jujuDB)
	c := db.C(metadataCollection)
	txnRunner := jujutxn.NewRunner(jujutxn.RunnerParams{Database: db})
	managedStorage := blobstore.NewManagedStorage(db, rs)
	storage := binarystorageNew(uuid, managedStorage, c, txnRunner)
	return &storageCloser{storage, session}
}

// ToolsStorage returns a new toolstorage.StorageCloser
// that stores tools metadata in the "juju" database''
// "toolsmetadata" collection.
//
// TODO(axw) remove this, add a constructor function in toolstorage.
func (st *State) ToolsStorage() (toolstorage.StorageCloser, error) {
	uuid := st.EnvironUUID()
	session := st.session.Copy()
	rs := blobstore.NewGridFS(blobstoreDB, uuid, session)
	db := session.DB(jujuDB)
	metadataCollection := db.C(toolsmetadataC)
	txnRunner := jujutxn.NewRunner(jujutxn.RunnerParams{Database: db})
	managedStorage := blobstore.NewManagedStorage(db, rs)
	storage := toolstorageNewStorage(uuid, managedStorage, metadataCollection, txnRunner)
	return &toolsStorageCloser{storage, session}, nil
}

func (s *ImageSuite) SetUpTest(c *gc.C) {
	s.BaseSuite.SetUpTest(c)
	s.mongo = &gitjujutesting.MgoInstance{}
	s.mongo.Start(nil)

	var err error
	s.session, err = s.mongo.Dial()
	c.Assert(err, gc.IsNil)
	s.storage = imagestorage.NewStorage(s.session, "my-uuid")
	s.metadataCollection = imagestorage.MetadataCollection(s.storage)
	s.txnRunner = jujutxn.NewRunner(jujutxn.RunnerParams{Database: s.metadataCollection.Database})
	s.patchTransactionRunner()
}

func (s *managedStorageSuite) SetUpTest(c *gc.C) {
	s.IsolationSuite.SetUpTest(c)
	s.MgoSuite.SetUpTest(c)
	s.db = s.Session.DB("blobstore")
	s.resourceStorage = blobstore.NewGridFS("storage", "test", s.Session)
	s.managedStorage = blobstore.NewManagedStorage(s.db, s.resourceStorage)

	// For testing, we need to ensure there's a single txnRunner for all operations.
	s.txnRunner = jujutxn.NewRunner(jujutxn.RunnerParams{Database: s.db})
	txnRunnerFunc := func(db *mgo.Database) jujutxn.Runner {
		return s.txnRunner
	}
	s.PatchValue(blobstore.TxnRunner, txnRunnerFunc)
}

func (s *txnSuite) TestRunFailureIntermittentUnexpectedMessage(c *gc.C) {
	runner := jujutxn.NewRunner(jujutxn.RunnerParams{})
	fake := &fakeRunner{errors: []error{errors.New("unexpected message")}}
	jujutxn.SetRunnerFunc(runner, fake.new)
	tries := 0
	// Doesn't matter what this returns as long as it isn't an error.
	buildTxn := func(attempt int) ([]txn.Op, error) {
		tries++
		return nil, nil
	}
	err := runner.Run(buildTxn)
	c.Check(err, gc.Equals, nil)
	c.Check(tries, gc.Equals, 2)
}

func (s *resourceCatalogSuite) SetUpTest(c *gc.C) {
	s.IsolationSuite.SetUpTest(c)
	s.MgoSuite.SetUpTest(c)
	db := s.Session.DB("blobstore")
	s.collection = db.C("storedResources")
	s.rCatalog = blobstore.NewResourceCatalog(db)

	// For testing, we need to ensure there's a single txnRunner for all operations.
	s.txnRunner = txn.NewRunner(txn.RunnerParams{Database: db})
	txnRunnerFunc := func(db *mgo.Database) txn.Runner {
		return s.txnRunner
	}
	s.PatchValue(blobstore.TxnRunner, txnRunnerFunc)
}

func (s *ToolsSuite) SetUpTest(c *gc.C) {
	s.BaseSuite.SetUpTest(c)
	s.mongo = &gitjujutesting.MgoInstance{}
	s.mongo.Start(nil)

	var err error
	s.session, err = s.mongo.Dial()
	c.Assert(err, jc.ErrorIsNil)
	rs := blobstore.NewGridFS("blobstore", "my-uuid", s.session)
	catalogue := s.session.DB("catalogue")
	s.managedStorage = blobstore.NewManagedStorage(catalogue, rs)
	s.metadataCollection = catalogue.C("toolsmetadata")
	s.txnRunner = jujutxn.NewRunner(jujutxn.RunnerParams{Database: catalogue})
	s.storage = toolstorage.NewStorage("my-uuid", s.managedStorage, s.metadataCollection, s.txnRunner)
}

// newStorageDBWrapper returns a DB operator for the , with its own session.
func newStorageDBWrapper(db *mgo.Database, metaColl, envUUID string) *storageDBWrapper {
	session := db.Session.Copy()
	db = db.With(session)

	coll := db.C(metaColl)
	txnRunner := jujutxn.NewRunner(jujutxn.RunnerParams{Database: db})
	dbWrap := storageDBWrapper{
		session:   session,
		db:        db,
		metaColl:  coll,
		txnRunner: txnRunner,
		envUUID:   envUUID,
	}
	return &dbWrap
}

// Copy returns a copy of the operator.
func (b *storageDBWrapper) Copy() *storageDBWrapper {
	session := b.session.Copy()

	coll := b.metaColl.With(session)
	db := coll.Database
	txnRunner := jujutxn.NewRunner(jujutxn.RunnerParams{Database: db})
	dbWrap := storageDBWrapper{
		session:   session,
		db:        db,
		metaColl:  coll,
		txnRunner: txnRunner,
		envUUID:   b.envUUID,
	}
	return &dbWrap
}

func (s *binaryStorageSuite) SetUpTest(c *gc.C) {
	s.BaseSuite.SetUpTest(c)
	s.mongo = &gitjujutesting.MgoInstance{}
	s.mongo.Start(nil)

	var err error
	var closer func()
	s.session, err = s.mongo.Dial()
	c.Assert(err, jc.ErrorIsNil)
	rs := blobstore.NewGridFS("blobstore", "blobstore", s.session)
	catalogue := s.session.DB("catalogue")
	s.managedStorage = blobstore.NewManagedStorage(catalogue, rs)
	s.metadataCollection, closer = mongo.CollectionFromName(catalogue, "binarymetadata")
	s.AddCleanup(func(*gc.C) { closer() })
	s.txnRunner = jujutxn.NewRunner(jujutxn.RunnerParams{Database: catalogue})
	s.storage = binarystorage.New("my-uuid", s.managedStorage, s.metadataCollection, s.txnRunner)
}

// TransactionRunner is part of the Database interface.
func (db *database) TransactionRunner() (runner jujutxn.Runner, closer SessionCloser) {
	runner = db.runner
	closer = dontCloseAnything
	if runner == nil {
		raw := db.raw
		if !db.ownSession {
			session := raw.Session.Copy()
			raw = raw.With(session)
			closer = session.Close
		}
		params := jujutxn.RunnerParams{Database: raw}
		runner = jujutxn.NewRunner(params)
	}
	return &multiModelRunner{
		rawRunner: runner,
		modelUUID: db.modelUUID,
		schema:    db.schema,
	}, closer
}

	return blobstore.NewManagedStorage(metadataDb, rs)
}

func (s *imageStorage) getManagedStorage(session *mgo.Session) blobstore.ManagedStorage {
	return getManagedStorage(session)
}

func (s *imageStorage) txnRunner(session *mgo.Session) jujutxn.Runner {
	db := s.metadataCollection.Database
	runnerDb := db.With(session)
	return txnRunner(runnerDb)
}

// Override for testing.
var txnRunner = func(db *mgo.Database) jujutxn.Runner {
	return jujutxn.NewRunner(jujutxn.RunnerParams{Database: db})
}

// AddImage is defined on the Storage interface.
func (s *imageStorage) AddImage(r io.Reader, metadata *Metadata) (resultErr error) {
	session := s.blobDb.Session.Copy()
	defer session.Close()
	managedStorage := s.getManagedStorage(session)
	path := imagePath(metadata.Kind, metadata.Series, metadata.Arch, metadata.SHA256)
	if err := managedStorage.PutForEnvironment(s.envUUID, path, r, metadata.Size); err != nil {
		return errors.Annotate(err, "cannot store image")
	}
	defer func() {
		if resultErr == nil {
			return
		}

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

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