package sources

import (
	"database/sql"
	"testing"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/stretchr/testify/assert"
	"github.com/tidbcloud/observability/data-collector/pkg/collector/config"
)

func TestOpenDB(t *testing.T) {
	cfg := config.DBSource{
		DBUrl: "test",
	}
	src := NewDBSource(cfg)
	assert.NotEmpty(t, src.GetConfig())
	_, err := OpenDBfn(cfg)
	assert.Error(t, err)
}
func TestDBRunSuccess(t *testing.T) {
	db, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal("sql mock init failed")
	}
	//inject mock db
	OpenDBfn = func(cfg config.DBSource) (*sql.DB, error) {
		return db, nil
	}
	defer db.Close()
	dbs := NewDBSource(config.DBSource{
		DBUrl: "test",
		Tables: []struct {
			Name    string           `yaml:"name" mapstructure:"name"`
			Cols    []config.ColType `yaml:"cols" mapstructure:"cols"`
			TimeCol string           `yaml:"time_col" mapstructure:"time_col"`
		}{
			{
				Name: "test",
				Cols: []config.ColType{
					{
						Name: "col1",
						Type: "string",
					},
					{
						Name: "col2",
						Type: "binary",
					},
				},
				TimeCol: "created_at",
			},
		},
	})
	// mock select count return 1 row
	mock.ExpectQuery("select count").WillReturnRows(sqlmock.NewRows([]string{"count(*)"}).AddRow(1))
	// mock select data return 1 row
	mock.ExpectQuery("select col1,col2").WillReturnRows(sqlmock.NewRows([]string{"col1", "col2"}).AddRow(1, []byte{30, 50, 20}))
	result, err := dbs.Run(1678283417, 1678283817)
	assert.NoError(t, err)
	assert.Equal(t, 1, len(result))
}

func TestDBRunFailed(t *testing.T) {
	db, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal("sql mock init failed")
	}
	dbs := NewDBSource(config.DBSource{
		DBUrl: "test",
		Tables: []struct {
			Name    string           `yaml:"name" mapstructure:"name"`
			Cols    []config.ColType `yaml:"cols" mapstructure:"cols"`
			TimeCol string           `yaml:"time_col" mapstructure:"time_col"`
		}{
			{
				Name: "test",
				Cols: []config.ColType{
					{
						Name: "col1",
						Type: "string",
					},
				},
				TimeCol: "created_at",
			},
		},
	})
	//open db failed
	OpenDBfn = func(cfg config.DBSource) (*sql.DB, error) {
		return nil, sql.ErrConnDone
	}
	_, err = dbs.Run(1678283417, 1678283817)
	assert.Error(t, err)
	//inject mock db
	OpenDBfn = func(cfg config.DBSource) (*sql.DB, error) {
		return db, nil
	}
	//mock query error
	mock.ExpectQuery("select count").WillReturnError(sql.ErrNoRows)
	result, err := dbs.Run(1678283417, 1678283817)
	assert.Error(t, err)
	assert.Equal(t, 0, len(result))
	db, mock, err = sqlmock.New()
	if err != nil {
		t.Fatal("sql mock init failed")
	}
	OpenDBfn = func(cfg config.DBSource) (*sql.DB, error) {
		return db, nil
	}
	//table count is empty
	mock.ExpectQuery("select count").WillReturnRows(sqlmock.NewRows([]string{"count(*)"}).AddRow(0))
	result, err = dbs.Run(1678283417, 1678283817)
	assert.NoError(t, err)
	assert.Equal(t, 0, len(result))

}