mirror of
https://github.com/activecm/rita
synced 2026-06-08 13:02:45 +00:00
Liza Tsibur
412 lines
15 KiB
Go
412 lines
15 KiB
Go
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")
|
|
})
|
|
}
|