Change db structure and add test cases

Signed-off-by: Qifan Deng <[email protected]>

Author: pancak3

pkg/llm-d-inference-sim/.llm-d/test.invalid.column.sqlite3

@@ -0,0 +1 @@
+Hello world!

pkg/llm-d-inference-sim/.llm-d/test.invalid.sqlite3

@@ -38,6 +38,17 @@ type Dataset struct {
 	logger logr.Logger
 }
 
+// use constants for expected column names and types
+const (
+	tableName         = "llmd"
+	promptHashCol     = "prompt_hash"
+	genTokensCol      = "gen_tokens"
+	nGenTokensCol     = "n_gen_tokens"
+	promptHashColType = "BLOB"
+	genTokensColType  = "JSON"
+	nGenTokensColType = "INTEGER"
+)
+
 func (d *Dataset) downloadDataset(url string, savePath string) error {
 	// Set up signal handling for Ctrl+C (SIGINT)
 	ctx, cancel := context.WithCancel(context.Background())
@@ -169,7 +180,73 @@ func (pr *progressReader) logProgress(pct int) {
 	}
 }
 
+func (d *Dataset) verifyDB() error {
+	rows, err := d.db.Query("PRAGMA table_info(" + tableName + ");")
+	if err != nil {
+		return fmt.Errorf("failed to query table info for `%s`: %w", tableName, err)
+	}
+	defer func() {
+		if cerr := rows.Close(); cerr != nil {
+			d.logger.Error(cerr, "failed to close rows after querying table info")
+		}
+	}()
+
+	expectedColumns := map[string]string{
+		promptHashCol: promptHashColType,
+		genTokensCol:  genTokensColType,
+		nGenTokensCol: nGenTokensColType,
+	}
+
+	columnsFound := make(map[string]bool)
+
+	var (
+		columnName string
+		columnType string
+		cid        int
+		notnull    int
+		dfltValue  interface{}
+		pk         int
+	)
+
+	for rows.Next() {
+		err := rows.Scan(&cid, &columnName, &columnType, &notnull, &dfltValue, &pk)
+		if err != nil {
+			return fmt.Errorf("failed to scan table info row: %w", err)
+		}
+		if expectedType, exists := expectedColumns[columnName]; exists {
+			if columnType != expectedType {
+				return fmt.Errorf("column %s has incorrect type: expected %s, got %s", columnName, expectedType, columnType)
+			}
+			columnsFound[columnName] = true
+		}
+	}
+
+	for col := range expectedColumns {
+		if !columnsFound[col] {
+			return fmt.Errorf("missing expected column in %s table: %s", tableName, col)
+		}
+	}
+
+	return nil
+}
+
+func (d *Dataset) getRecordsCount() (int, error) {
+	var count int
+	err := d.db.QueryRow("SELECT COUNT(" + promptHashCol + ") FROM " + tableName + ";").Scan(&count)
+	if err != nil {
+		return 0, fmt.Errorf("failed to query database: %w", err)
+	}
+	return count, nil
+}
+
 func (d *Dataset) connectToDB(path string) error {
+	if d.db != nil {
+		err := d.db.Close()
+		if err != nil {
+			d.logger.Error(err, "failed to close existing database connection")
+		}
+		d.db = nil
+	}
 	// check if file exists
 	_, err := os.Stat(path)
 	if err != nil {
@@ -180,13 +257,15 @@ func (d *Dataset) connectToDB(path string) error {
 		return fmt.Errorf("failed to open database: %w", err)
 	}
 
-	var count int
-	err = d.db.QueryRow("SELECT COUNT(generated) FROM llmd;").Scan(&count)
+	err = d.verifyDB()
+
 	if err != nil {
-		err := d.db.Close()
-		if err != nil {
-			d.logger.Error(err, "failed to close database after query failure")
-		}
+		return fmt.Errorf("failed to verify database: %w", err)
+	}
+
+	count, err := d.getRecordsCount()
+	if err != nil {
+		d.logger.Error(err, "failed to get records count")
 		return fmt.Errorf("failed to query database: %w", err)
 	}
 	d.logger.Info("Database connected successfully", "path", path, "records count", count)

pkg/llm-d-inference-sim/.llm-d/test.invalid.table.sqlite3

@@ -17,6 +17,7 @@ limitations under the License.
 package llmdinferencesim
 
 import (
+	"fmt"
 	"os"
 
 	"github.com/go-logr/logr"
@@ -28,19 +29,31 @@ import (
 
 var _ = Describe("Dataset", func() {
 	var (
-		dataset     *Dataset
-		file_folder string
-		savePath    string
+		dataset               *Dataset
+		file_folder           string
+		savePath              string
+		validDBPath           string
+		pathToInvalidDB       string
+		pathNotExist          string
+		pathToInvalidTableDB  string
+		pathToInvalidColumnDB string
+		pathToInvalidTypeDB   string
 	)
 
 	BeforeEach(func() {
 		dataset = &Dataset{
 			logger: logr.Discard(),
 		}
-		file_folder = "./.llm-d"
+		file_folder = ".llm-d"
 		savePath = file_folder + "/test.sqlite3"
 		err := os.MkdirAll(file_folder, os.ModePerm)
 		Expect(err).NotTo(HaveOccurred())
+		validDBPath = file_folder + "/test.valid.sqlite3"
+		pathNotExist = file_folder + "/test.notexist.sqlite3"
+		pathToInvalidDB = file_folder + "/test.invalid.sqlite3"
+		pathToInvalidTableDB = file_folder + "/test.invalid.table.sqlite3"
+		pathToInvalidColumnDB = file_folder + "/test.invalid.column.sqlite3"
+		pathToInvalidTypeDB = file_folder + "/test.invalid.type.sqlite3"
 	})
 
 	AfterEach(func() {
@@ -72,24 +85,47 @@ var _ = Describe("Dataset", func() {
 	})
 
 	It("should successfully init dataset", func() {
-		validDBPath := file_folder + "/test.valid.sqlite3"
 		err := dataset.Init(validDBPath, "", "")
+		// debug: get the realpath
+		wd, _ := os.Getwd()
+		realpath := fmt.Sprintf("%s/%s", wd, validDBPath)
+		fmt.Println("Using realpath:", realpath)
 		Expect(err).NotTo(HaveOccurred())
 
-		row := dataset.db.QueryRow("SELECT generated FROM llmd WHERE prompt_hash=X'b94d27b9934d041c52e5b721d7373f13a07ed5e79179d63c5d8a0c102a9d00b2';")
-		var value string
-		err = row.Scan(&value)
+		row := dataset.db.QueryRow("SELECT n_gen_tokens FROM llmd WHERE prompt_hash=X'b94d27b9934d041c52e5b721d7373f13a07ed5e79179d63c5d8a0c102a9d00b2';")
+		var n_gen_tokens int
+		err = row.Scan(&n_gen_tokens)
 		Expect(err).NotTo(HaveOccurred())
-		Expect(value).To(Equal("world!"))
+		Expect(n_gen_tokens).To(Equal(3))
 	})
 
-	It("should raise err with invalid DB content", func() {
-		err := dataset.connectToDB(file_folder)
+	It("should return error for non-existing DB path", func() {
+		err := dataset.connectToDB(pathNotExist)
 		Expect(err).To(HaveOccurred())
-		// read from the db to verify it's not valid
-		row := dataset.db.QueryRow("SELECT * FROM llmd;")
-		var value string
-		err = row.Scan(&value)
+		Expect(err.Error()).To(ContainSubstring("database file does not exist"))
+	})
+	
+	It("should return error for invalid DB file", func() {
+		err := dataset.connectToDB(pathToInvalidDB)
+		Expect(err).To(HaveOccurred())
+		Expect(err.Error()).To(ContainSubstring("file is not a database"))
+	})
+	
+	It("should return error for DB with invalid table", func() {
+		err := dataset.connectToDB(pathToInvalidTableDB)
+		Expect(err).To(HaveOccurred())
+		Expect(err.Error()).To(ContainSubstring("failed to verify database"))
+	})
+
+	It("should return error for DB with invalid column", func() {
+		err := dataset.connectToDB(pathToInvalidColumnDB)
+		Expect(err).To(HaveOccurred())
+		Expect(err.Error()).To(ContainSubstring("missing expected column"))
+	})
+	
+	It("should return error for DB with invalid column type", func() {
+		err := dataset.connectToDB(pathToInvalidTypeDB)
 		Expect(err).To(HaveOccurred())
+		Expect(err.Error()).To(ContainSubstring("incorrect type"))
 	})
 })

pkg/llm-d-inference-sim/.llm-d/test.invalid.type.sqlite3

@@ -221,7 +221,7 @@ func (s *VllmSimulator) startSim(ctx context.Context) error {
 		s.logger.Info("No dataset provided, will generate random responses")
 	} else {
 		dataset := &Dataset{
-			logger:   s.logger,
+			logger: s.logger,
 		}
 		err = dataset.Init(s.config.Dataset.Path, s.config.Dataset.Url, s.config.Dataset.SavePath)
 		if err != nil {