github.com/juju/juju@v0.0.0-20240430160146-1752b71fcf00/database/testing/simplesuite.go (about)

     1  // Copyright 2022 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package testing
     5  
     6  import (
     7  	"database/sql"
     8  	"fmt"
     9  
    10  	"github.com/juju/testing"
    11  	jc "github.com/juju/testing/checkers"
    12  	_ "github.com/mattn/go-sqlite3"
    13  	gc "gopkg.in/check.v1"
    14  
    15  	coredatabase "github.com/juju/juju/core/database"
    16  	"github.com/juju/juju/database/schema"
    17  )
    18  
    19  // ControllerSuite is used to provide an in-memory sql.DB reference to tests.
    20  // It is pre-populated with the controller schema.
    21  type ControllerSuite struct {
    22  	testing.IsolationSuite
    23  
    24  	db        *sql.DB
    25  	trackedDB coredatabase.TrackedDB
    26  }
    27  
    28  // SetUpTest creates a new sql.DB reference and ensures that the
    29  // controller schema is applied successfully.
    30  func (s *ControllerSuite) SetUpTest(c *gc.C) {
    31  	s.IsolationSuite.SetUpTest(c)
    32  
    33  	// Do not be tempted in moving to :memory: mode for this test suite. It will
    34  	// fail in non-deterministic ways. Unfortunately :memory: mode is not
    35  	// completely goroutine safe.
    36  	s.db = s.NewCleanDB(c)
    37  
    38  	s.trackedDB = &trackedDB{
    39  		db: s.db,
    40  	}
    41  
    42  	s.ApplyControllerDDL(c, s.db)
    43  }
    44  
    45  func (s *ControllerSuite) TearDownTest(c *gc.C) {
    46  	if s.db != nil {
    47  		c.Logf("Closing DB")
    48  		err := s.db.Close()
    49  		c.Assert(err, jc.ErrorIsNil)
    50  	}
    51  
    52  	s.IsolationSuite.TearDownTest(c)
    53  }
    54  
    55  // DB returns a sql.DB reference.
    56  func (s *ControllerSuite) DB() *sql.DB {
    57  	return s.db
    58  }
    59  
    60  // TrackedDB returns a TrackedDB reference.
    61  func (s *ControllerSuite) TrackedDB() coredatabase.TrackedDB {
    62  	return s.trackedDB
    63  }
    64  
    65  // NewCleanDB returns a new sql.DB reference.
    66  func (s *ControllerSuite) NewCleanDB(c *gc.C) *sql.DB {
    67  	dir := c.MkDir()
    68  
    69  	url := fmt.Sprintf("file:%s/db.sqlite3?_foreign_keys=1", dir)
    70  	c.Logf("Opening sqlite3 db with: %v", url)
    71  
    72  	db, err := sql.Open("sqlite3", url)
    73  	c.Assert(err, jc.ErrorIsNil)
    74  
    75  	return db
    76  }
    77  
    78  // ApplyControllerDDL applies the controller schema to the provided sql.DB.
    79  // This is useful for tests that need to apply the schema to a new DB.
    80  func (s *ControllerSuite) ApplyControllerDDL(c *gc.C, db *sql.DB) {
    81  	tx, err := s.db.Begin()
    82  	c.Assert(err, jc.ErrorIsNil)
    83  
    84  	for idx, stmt := range schema.ControllerDDL() {
    85  		c.Logf("Executing schema DDL index: %v", idx)
    86  		_, err := tx.Exec(stmt)
    87  		c.Assert(err, jc.ErrorIsNil)
    88  	}
    89  
    90  	c.Logf("Committing schema DDL")
    91  	err = tx.Commit()
    92  	c.Assert(err, jc.ErrorIsNil)
    93  }