package database_test import ( "context" "fmt" "log" "os" "path/filepath" "testing" "time" "github.com/activecm/rita/v5/cmd" "github.com/activecm/rita/v5/config" "github.com/activecm/rita/v5/database" "github.com/ClickHouse/clickhouse-go/v2" "github.com/joho/godotenv" "github.com/spf13/afero" "github.com/stretchr/testify/require" "github.com/stretchr/testify/suite" "github.com/testcontainers/testcontainers-go" cl "github.com/testcontainers/testcontainers-go/modules/clickhouse" ) const ConfigPath = "../integration/test_config.hjson" const TestDataPath = "../test_data" type DatabaseTestSuite struct { suite.Suite cfg *config.Config clickhouseContainer *cl.ClickHouseContainer clickhouseConnection string server *database.ServerConn } func TestMain(m *testing.M) { // load environment variables with panic prevention if err := godotenv.Overload("../.env", "../integration/test.env"); err != nil { log.Fatalf("Error loading .env file: %v", err) } os.Exit(m.Run()) } func TestDatabase(t *testing.T) { suite.Run(t, new(DatabaseTestSuite)) } // SetupSuite is run once before the first test starts func (d *DatabaseTestSuite) SetupSuite() { t := d.T() // load the config file cfg, err := config.ReadFileConfig(afero.NewOsFs(), ConfigPath) require.NoError(t, err, "config should load without error") // start clickhouse container d.SetupClickHouse(t) // update the config to use the clickhouse container connection cfg.Env.DBConnection = d.clickhouseConnection d.cfg = cfg // connect to clickhouse server server, err := database.ConnectToServer(context.Background(), d.cfg) require.NoError(t, err, "connecting to server should not produce an error") d.server = server } // TearDownSuite is run once after all tests have finished func (d *DatabaseTestSuite) TearDownSuite() { if err := d.clickhouseContainer.Terminate(context.Background()); err != nil { log.Fatalf("failed to terminate clickhouse container: %s", err) } } // SetupTest is run before each test method // func (d *DatabaseTestSuite) SetupTest() {} // TearDownTest is run after each test method // func (d *DatabaseTestSuite) TearDownTest() {} // SetupSubTest is run before each subtest func (d *DatabaseTestSuite) SetupSubTest() { t := d.T() // fmt.Println("Running setup subtest...") // drop all databases that may have been created during subtest if d.server != nil && d.server.Conn != nil { dbs, err := d.server.ListImportDatabases() require.NoError(t, err, "listing databases should not produce an error") for _, db := range dbs { err := d.server.DeleteSensorDB(db.Name) require.NoError(t, err, "dropping database should not produce an error") } } } // TearDownSubTest is run after each subtest // func (d *DatabaseTestSuite) TearDownSubTest() {} // SetupClickHouse creates a ClickHouse container using the test.docker-compose.yml and handles taking it down when complete func (d *DatabaseTestSuite) SetupClickHouse(t *testing.T) { t.Helper() // get ClickHouse version from environment version := os.Getenv("CLICKHOUSE_VERSION") require.NotEmpty(t, version, "CLICKHOUSE_VERSION environment variable must be set") // create ClickHouse container ctx := context.Background() clickHouseContainer, err := cl.RunContainer(ctx, testcontainers.WithImage(fmt.Sprintf("clickhouse/clickhouse-server:%s-alpine", version)), cl.WithUsername("default"), cl.WithPassword(""), cl.WithDatabase("default"), cl.WithConfigFile(filepath.Join("../deployment/", "config.xml")), ) require.NoError(t, err, "failed to start clickHouse container") // get connection host connectionHost, err := clickHouseContainer.ConnectionHost(ctx) require.NoError(t, err, "failed to get clickHouse connection host") // set container and connection host d.clickhouseContainer = clickHouseContainer d.clickhouseConnection = connectionHost } func (d *DatabaseTestSuite) TestConnectToServer() { // connect to the server with valid configuration d.Run("Successful Server Connection", func() { t := d.T() server, err := database.ConnectToServer(context.Background(), d.cfg) require.NoError(t, err, "connecting to clickhouse server should not produce an error") require.NotNil(t, server, "server connection object should not be nil") // ping to ensure the connection is valid err = server.Conn.Ping(context.Background()) require.NoError(t, err, "pinging clickhouse server should not produce an error") }) // attempt to connect with invalid configuration d.Run("Failed Server Connection", func() { t := d.T() invalidCfg := *d.cfg invalidCfg.Env.DBConnection = "invalid connection string" server, err := database.ConnectToServer(context.Background(), &invalidCfg) require.Error(t, err, "connecting with invalid configuration should produce an error") require.Nil(t, server, "server connection object should be nil on failed connection") }) } func (d *DatabaseTestSuite) TestDeleteSensorDB() { d.Run("Drop Existing Database", func() { t := d.T() _, err := cmd.RunImportCmd(time.Now(), d.cfg, afero.NewOsFs(), "../test_data/valid_tsv", "testDB", false, false) require.NoError(t, err, "importing data should not produce an error") db, err := database.ConnectToDB(context.Background(), "testDB", d.cfg, nil) require.NoError(t, err, "connecting to created database should not produce an error") require.NotNil(t, db) // drop the database err = d.server.DeleteSensorDB("testDB") require.NoError(t, err, "dropping database should not produce an error") d.checkDatabaseDeletion("testDB") }) d.Run("Drop Non-Existent Database", func() { t := d.T() // attempt to drop a database that doesn't exist err := d.server.DeleteSensorDB("nonExistentDB") require.NoError(t, err, "attempting to drop a non-existent database should not produce an error") }) } func (d *DatabaseTestSuite) checkDatabaseDeletion(dbName string) { t := d.T() // attempt to connect to the dropped database db, err := database.ConnectToDB(context.Background(), dbName, d.cfg, nil) require.Error(t, err, "connecting to a dropped database should produce an error") require.Nil(t, db) ctx := d.server.QueryParameters(clickhouse.Parameters{ "database": dbName, }) // check for db in min_max var count uint64 err = d.server.Conn.QueryRow(ctx, ` SELECT count() FROM metadatabase.min_max WHERE database = {database:String} `).Scan(&count) require.NoError(t, err, "querying metadatabase.min_max should not produce an error") require.EqualValues(t, 0, count, "there should be no records for a deleted dataset in metadatabase.min_max, database: %s", dbName) // check for db in files err = d.server.Conn.QueryRow(ctx, ` SELECT count() FROM metadatabase.files WHERE database = {database:String} `).Scan(&count) require.NoError(t, err, "querying metadatabase.files should not produce an error") require.EqualValues(t, 0, count, "there should be no records for a deleted dataset in metadatabase.files, database: %s", dbName) } func (d *DatabaseTestSuite) checkDatabaseNonDeletion(dbName string) { t := d.T() // attempt to connect to the database db, err := database.ConnectToDB(context.Background(), dbName, d.cfg, nil) require.NoError(t, err, "connecting to a database that was not dropped should not produce an error") require.NotNil(t, db) ctx := d.server.QueryParameters(clickhouse.Parameters{ "database": dbName, }) // check for db in min_max var count uint64 err = d.server.Conn.QueryRow(ctx, ` SELECT count() FROM metadatabase.min_max WHERE database = {database:String} `).Scan(&count) require.NoError(t, err, "querying metadatabase.min_max should not produce an error") require.Greater(t, count, uint64(0), "there should be at least 1 record for a dataset in metadatabase.min_max, database: %s", dbName) // check for db in files err = d.server.Conn.QueryRow(ctx, ` SELECT count() FROM metadatabase.files WHERE database = {database:String} `).Scan(&count) require.NoError(t, err, "querying metadatabase.files should not produce an error") require.Greater(t, count, uint64(0), "there should be at least 1 record for a dataset in metadatabase.files, database: %s", dbName) } func (d *DatabaseTestSuite) TestDropMultipleSensorDatabases() { // helper function to create databases with various prefixes and suffixes databases := []string{"bingbong", "prefix_bingbong", "bingbong123", "prefix_bingbong123"} createDatabases := func(t *testing.T) { t.Helper() for _, dbName := range databases { // create a database name with no prefix or suffix _, err := cmd.RunImportCmd(time.Now(), d.cfg, afero.NewOsFs(), "../test_data/valid_tsv", dbName, false, false) require.NoError(t, err, "importing data should not produce an error") } } d.Run("Drop Databases with Prefix Wildcard", func() { t := d.T() // create the databases createDatabases(t) deleted := databases[:2] notDeleted := databases[2:] // drop databases with a prefix wildcard numDeleted, err := d.server.DropMultipleSensorDatabases("bingbong", true, false) require.NoError(t, err, "dropping databases with prefix wildcard should not produce an error") require.Equal(t, 2, numDeleted, "should delete exactly 2 databases") // should match the first two databases created for _, dbName := range deleted { d.checkDatabaseDeletion(dbName) } for _, dbName := range notDeleted { d.checkDatabaseNonDeletion(dbName) } }) d.Run("Drop Databases with Suffix Wildcard", func() { t := d.T() // create the databases createDatabases(t) deleted := []string{databases[0], databases[2]} notDeleted := []string{databases[1], databases[3]} // drop databases with a suffix wildcard numDeleted, err := d.server.DropMultipleSensorDatabases("bingbong", false, true) require.NoError(t, err, "dropping databases with suffix wildcard should not produce an error") require.Equal(t, 2, numDeleted, "should delete exactly 2 databases") // should match the first and third databases created for _, dbName := range deleted { d.checkDatabaseDeletion(dbName) } for _, dbName := range notDeleted { d.checkDatabaseNonDeletion(dbName) } }) d.Run("Drop Databases with Both Wildcards", func() { t := d.T() // create the databases createDatabases(t) // drop databases with both wildcards numDeleted, err := d.server.DropMultipleSensorDatabases("bingbong", true, true) require.NoError(t, err, "dropping databases with both wildcards should not produce an error") require.Equal(t, 4, numDeleted, "should delete all databases matching the wildcard pattern") // should match all databases created for _, dbName := range databases { d.checkDatabaseDeletion(dbName) } }) d.Run("Drop Databases with No Wildcards", func() { t := d.T() // create the databases createDatabases(t) // drop databases without wildcards numDeleted, err := d.server.DropMultipleSensorDatabases("bingbong", false, false) require.Error(t, err, "dropping databases without specifying a wildcard should produce an error") require.Equal(t, 0, numDeleted, "no databases should be deleted if no wildcard is specified") for _, dbName := range databases { d.checkDatabaseNonDeletion(dbName) } }) } func (d *DatabaseTestSuite) TestListImportDatabases() { d.Run("List Databases", func() { t := d.T() _, err := cmd.RunImportCmd(time.Now(), d.cfg, afero.NewOsFs(), "../test_data/valid_tsv", "testDB", false, true) require.NoError(t, err, "importing data should not produce an error") _, err = cmd.RunImportCmd(time.Now(), d.cfg, afero.NewOsFs(), "../test_data/valid_tsv", "testDB2", false, true) require.NoError(t, err, "importing data should not produce an error") dbs, err := d.server.ListImportDatabases() require.NoError(t, err, "listing databases should not produce an error") require.Len(t, dbs, 2, "two databases should be listed") dbString := database.GetFlatDatabaseList(dbs) require.Containsf(t, dbString, "testDB", "testDB should be listed") require.Contains(t, dbString, "testDB2", "testDB2 should be listed") }) d.Run("No Databases", func() { t := d.T() dbs, err := d.server.ListImportDatabases() require.NoError(t, err, "listing databases should not produce an error") require.Nil(t, dbs, "databases should be nil") }) d.Run("Missing Metadatabase", func() { t := d.T() _, err := cmd.RunImportCmd(time.Now(), d.cfg, afero.NewOsFs(), "../test_data/valid_tsv", "testDB", false, false) require.NoError(t, err, "importing data should not produce an error") err = d.server.Conn.Exec(context.Background(), "DROP DATABASE IF EXISTS metadatabase") require.NoError(t, err, "dropping metadatabase should not produce an error") dbs, err := d.server.ListImportDatabases() require.NoError(t, err, "listing databases with missing metadatabase should not produce an error") require.Nil(t, dbs, "databases should be nil") }) } func (d *DatabaseTestSuite) TestGetRollingStatus() { d.Run("Get Status of Rolling Database", func() { t := d.T() _, err := cmd.RunImportCmd(time.Now(), d.cfg, afero.NewOsFs(), "../test_data/valid_tsv", "testDB", true, false) require.NoError(t, err, "importing data should not produce an error") status, err := database.GetRollingStatus(context.Background(), d.server.Conn, "testDB") require.NoError(t, err, "getting status of rolling database should not produce an error") fmt.Print("status: ", status) require.True(t, status, "status of rolling database should be true") }) d.Run("Get Status of Non-Rolling Database", func() { t := d.T() _, err := cmd.RunImportCmd(time.Now(), d.cfg, afero.NewOsFs(), "../test_data/valid_tsv", "testDB", false, false) require.NoError(t, err, "importing data should not produce an error") status, err := database.GetRollingStatus(context.Background(), d.server.Conn, "testDB") require.NoError(t, err, "getting status of non-rolling database should not produce an error") require.False(t, status, "status of non-rolling database should be false") }) d.Run("Get Status of Non-Existent Database", func() { t := d.T() status, err := database.GetRollingStatus(context.Background(), d.server.Conn, "testDB") require.Error(t, err, "getting status of non-existent database should produce an error") require.Equal(t, err, database.ErrDatabaseNotFound, "error should be database not found") require.False(t, status, "status of non-existent database should be false") }) } func (d *DatabaseTestSuite) TestDatabaseExists() { d.Run("Database Exists", func() { t := d.T() _, err := cmd.RunImportCmd(time.Now(), d.cfg, afero.NewOsFs(), "../test_data/valid_tsv", "testDB", false, false) require.NoError(t, err, "importing data should not produce an error") exists, err := database.DatabaseExists(context.Background(), d.server.Conn, "testDB") require.NoError(t, err, "checking if database exists should not produce an error") require.True(t, exists, "database should exist") }) d.Run("Database Does Not Exist", func() { t := d.T() exists, err := database.DatabaseExists(context.Background(), d.server.Conn, "testDB") require.NoError(t, err, "checking if database exists should not produce an error") require.False(t, exists, "database should not exist") }) }