database_test.go
1 // Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package database 16 17 import ( 18 "os" 19 "testing" 20 21 "github.com/stretchr/testify/assert" 22 "github.com/stretchr/testify/require" 23 ) 24 25 // newTestDB creates an in-memory SQLite DB for testing. 26 func newTestDB(t *testing.T) (*TaskStore, *ModelStore, func()) { 27 t.Helper() 28 29 // Use a temp file-based DB (pure in-memory ":memory:" doesn't work well 30 // with glebarez/sqlite + shared cache, so use a temp file instead). 31 f, err := os.CreateTemp("", "testdb-*.db") 32 require.NoError(t, err) 33 dbPath := f.Name() 34 f.Close() 35 36 cfg := NewConfig(dbPath) 37 db, err := InitDB(cfg) 38 require.NoError(t, err) 39 40 ts := NewTaskStore(db) 41 require.NoError(t, ts.Init()) 42 43 ms := NewModelStore(db) 44 require.NoError(t, ms.Init()) 45 46 cleanup := func() { 47 sqlDB, _ := db.DB() 48 if sqlDB != nil { 49 sqlDB.Close() 50 } 51 os.Remove(dbPath) 52 } 53 return ts, ms, cleanup 54 } 55 56 // --------------------------------------------------------------------------- 57 // InitDB / Config 58 // --------------------------------------------------------------------------- 59 60 func TestInitDB_InMemory(t *testing.T) { 61 f, err := os.CreateTemp("", "testdb-init-*.db") 62 require.NoError(t, err) 63 dbPath := f.Name() 64 f.Close() 65 defer os.Remove(dbPath) 66 67 cfg := NewConfig(dbPath) 68 assert.Equal(t, dbPath, cfg.DBPath) 69 70 db, err := InitDB(cfg) 71 require.NoError(t, err) 72 assert.NotNil(t, db) 73 74 sqlDB, err := db.DB() 75 require.NoError(t, err) 76 assert.NoError(t, sqlDB.Ping()) 77 sqlDB.Close() 78 } 79 80 func TestNewConfig(t *testing.T) { 81 cfg := NewConfig("/tmp/test.db") 82 assert.Equal(t, "/tmp/test.db", cfg.DBPath) 83 } 84 85 func TestLoadConfigFromEnv_Default(t *testing.T) { 86 os.Unsetenv("DB_PATH") 87 cfg := LoadConfigFromEnv() 88 assert.Equal(t, "db/tasks.db", cfg.DBPath) 89 } 90 91 func TestLoadConfigFromEnv_Override(t *testing.T) { 92 os.Setenv("DB_PATH", "/tmp/override.db") 93 defer os.Unsetenv("DB_PATH") 94 cfg := LoadConfigFromEnv() 95 assert.Equal(t, "/tmp/override.db", cfg.DBPath) 96 } 97 98 // --------------------------------------------------------------------------- 99 // TaskStore.Init 100 // --------------------------------------------------------------------------- 101 102 func TestTaskStore_Init(t *testing.T) { 103 ts, _, cleanup := newTestDB(t) 104 defer cleanup() 105 // Calling Init twice should be idempotent 106 assert.NoError(t, ts.Init()) 107 } 108 109 // --------------------------------------------------------------------------- 110 // TaskStore – User operations 111 // --------------------------------------------------------------------------- 112 113 func TestTaskStore_CreateAndGetUser(t *testing.T) { 114 ts, _, cleanup := newTestDB(t) 115 defer cleanup() 116 117 u := &User{ 118 UserID: "u-001", 119 Username: "alice", 120 Email: "alice@example.com", 121 IsActive: true, 122 } 123 require.NoError(t, ts.CreateUser(u)) 124 assert.Greater(t, u.CreatedAt, int64(0)) 125 126 got, err := ts.GetUser("alice") 127 require.NoError(t, err) 128 assert.Equal(t, "alice", got.Username) 129 assert.Equal(t, "alice@example.com", got.Email) 130 } 131 132 func TestTaskStore_GetUserByEmail(t *testing.T) { 133 ts, _, cleanup := newTestDB(t) 134 defer cleanup() 135 136 u := &User{UserID: "u-002", Username: "bob", Email: "bob@example.com"} 137 require.NoError(t, ts.CreateUser(u)) 138 139 got, err := ts.GetUserByEmail("bob@example.com") 140 require.NoError(t, err) 141 assert.Equal(t, "bob", got.Username) 142 } 143 144 func TestTaskStore_GetUser_NotFound(t *testing.T) { 145 ts, _, cleanup := newTestDB(t) 146 defer cleanup() 147 148 _, err := ts.GetUser("nonexistent") 149 assert.Error(t, err) 150 } 151 152 func TestTaskStore_CheckUserExists(t *testing.T) { 153 ts, _, cleanup := newTestDB(t) 154 defer cleanup() 155 156 u := &User{UserID: "u-003", Username: "charlie", Email: "charlie@example.com"} 157 require.NoError(t, ts.CreateUser(u)) 158 159 exists, err := ts.CheckUserExists("charlie@example.com") 160 require.NoError(t, err) 161 assert.True(t, exists) 162 163 exists, err = ts.CheckUserExists("nobody@example.com") 164 require.NoError(t, err) 165 assert.False(t, exists) 166 } 167 168 func TestTaskStore_DeleteUser(t *testing.T) { 169 ts, _, cleanup := newTestDB(t) 170 defer cleanup() 171 172 u := &User{UserID: "u-004", Username: "dave", Email: "dave@example.com"} 173 require.NoError(t, ts.CreateUser(u)) 174 require.NoError(t, ts.DeleteUser("dave@example.com")) 175 176 exists, err := ts.CheckUserExists("dave@example.com") 177 require.NoError(t, err) 178 assert.False(t, exists) 179 } 180 181 func TestTaskStore_UpdateUserFirstLogin(t *testing.T) { 182 ts, _, cleanup := newTestDB(t) 183 defer cleanup() 184 185 u := &User{UserID: "u-005", Username: "eve", Email: "eve@example.com", FirstLogin: true} 186 require.NoError(t, ts.CreateUser(u)) 187 188 require.NoError(t, ts.UpdateUserFirstLogin("eve", false)) 189 190 got, err := ts.GetUser("eve") 191 require.NoError(t, err) 192 assert.False(t, got.FirstLogin) 193 } 194 195 // --------------------------------------------------------------------------- 196 // TaskStore – Session (CreateSession / GetSession / UpdateSession / ListSessions) 197 // --------------------------------------------------------------------------- 198 199 func newTestSession(id, username, taskType string) *Session { 200 return &Session{ 201 ID: id, 202 Username: username, 203 Title: "Test session " + id, 204 TaskType: taskType, 205 Content: "test content", 206 Status: "todo", 207 } 208 } 209 210 func TestTaskStore_CreateSession(t *testing.T) { 211 ts, _, cleanup := newTestDB(t) 212 defer cleanup() 213 214 // Need a user first (foreign key) 215 require.NoError(t, ts.CreateUser(&User{UserID: "u1", Username: "testuser", Email: "t@t.com"})) 216 217 s := newTestSession("sess-001", "testuser", "ai_infra_scan") 218 require.NoError(t, ts.CreateSession(s)) 219 assert.Greater(t, s.CreatedAt, int64(0)) 220 assert.Greater(t, s.UpdatedAt, int64(0)) 221 } 222 223 func TestTaskStore_GetSession(t *testing.T) { 224 ts, _, cleanup := newTestDB(t) 225 defer cleanup() 226 227 require.NoError(t, ts.CreateUser(&User{UserID: "u2", Username: "user2", Email: "u2@t.com"})) 228 s := newTestSession("sess-002", "user2", "mcp_scan") 229 require.NoError(t, ts.CreateSession(s)) 230 231 got, err := ts.GetSession("sess-002") 232 require.NoError(t, err) 233 assert.Equal(t, "sess-002", got.ID) 234 assert.Equal(t, "mcp_scan", got.TaskType) 235 assert.Equal(t, "todo", got.Status) 236 } 237 238 func TestTaskStore_GetSession_NotFound(t *testing.T) { 239 ts, _, cleanup := newTestDB(t) 240 defer cleanup() 241 242 _, err := ts.GetSession("does-not-exist") 243 assert.Error(t, err) 244 } 245 246 func TestTaskStore_UpdateSessionStatus(t *testing.T) { 247 ts, _, cleanup := newTestDB(t) 248 defer cleanup() 249 250 require.NoError(t, ts.CreateUser(&User{UserID: "u3", Username: "user3", Email: "u3@t.com"})) 251 s := newTestSession("sess-003", "user3", "ai_infra_scan") 252 require.NoError(t, ts.CreateSession(s)) 253 254 require.NoError(t, ts.UpdateSessionStatus("sess-003", "doing")) 255 got, err := ts.GetSession("sess-003") 256 require.NoError(t, err) 257 assert.Equal(t, "doing", got.Status) 258 assert.NotNil(t, got.StartedAt) 259 260 require.NoError(t, ts.UpdateSessionStatus("sess-003", "done")) 261 got, err = ts.GetSession("sess-003") 262 require.NoError(t, err) 263 assert.Equal(t, "done", got.Status) 264 assert.NotNil(t, got.CompletedAt) 265 } 266 267 func TestTaskStore_UpdateSession(t *testing.T) { 268 ts, _, cleanup := newTestDB(t) 269 defer cleanup() 270 271 require.NoError(t, ts.CreateUser(&User{UserID: "u4", Username: "user4", Email: "u4@t.com"})) 272 s := newTestSession("sess-004", "user4", "ai_infra_scan") 273 require.NoError(t, ts.CreateSession(s)) 274 275 require.NoError(t, ts.UpdateSession("sess-004", map[string]interface{}{"title": "Updated Title"})) 276 got, err := ts.GetSession("sess-004") 277 require.NoError(t, err) 278 assert.Equal(t, "Updated Title", got.Title) 279 } 280 281 func TestTaskStore_DeleteSession(t *testing.T) { 282 ts, _, cleanup := newTestDB(t) 283 defer cleanup() 284 285 require.NoError(t, ts.CreateUser(&User{UserID: "u5", Username: "user5", Email: "u5@t.com"})) 286 s := newTestSession("sess-005", "user5", "ai_infra_scan") 287 require.NoError(t, ts.CreateSession(s)) 288 289 require.NoError(t, ts.DeleteSession("sess-005")) 290 _, err := ts.GetSession("sess-005") 291 assert.Error(t, err) 292 } 293 294 func TestTaskStore_ListSessions(t *testing.T) { 295 ts, _, cleanup := newTestDB(t) 296 defer cleanup() 297 298 require.NoError(t, ts.CreateUser(&User{UserID: "u6", Username: "user6", Email: "u6@t.com"})) 299 for i, id := range []string{"sess-a", "sess-b", "sess-c"} { 300 s := newTestSession(id, "user6", "mcp_scan") 301 s.Title = "Session " + string(rune('A'+i)) 302 require.NoError(t, ts.CreateSession(s)) 303 } 304 305 sessions, err := ts.GetUserSessions("user6") 306 require.NoError(t, err) 307 assert.Len(t, sessions, 3) 308 } 309 310 func TestTaskStore_GetUserSessionsByType(t *testing.T) { 311 ts, _, cleanup := newTestDB(t) 312 defer cleanup() 313 314 require.NoError(t, ts.CreateUser(&User{UserID: "u7", Username: "user7", Email: "u7@t.com"})) 315 require.NoError(t, ts.CreateSession(newTestSession("s7-1", "user7", "mcp_scan"))) 316 require.NoError(t, ts.CreateSession(newTestSession("s7-2", "user7", "ai_infra_scan"))) 317 require.NoError(t, ts.CreateSession(newTestSession("s7-3", "user7", "mcp_scan"))) 318 319 mcpSessions, err := ts.GetUserSessionsByType("user7", "mcp_scan") 320 require.NoError(t, err) 321 assert.Len(t, mcpSessions, 2) 322 323 allSessions, err := ts.GetUserSessionsByType("user7", "") 324 require.NoError(t, err) 325 assert.Len(t, allSessions, 3) 326 } 327 328 func TestTaskStore_SetShare(t *testing.T) { 329 ts, _, cleanup := newTestDB(t) 330 defer cleanup() 331 332 require.NoError(t, ts.CreateUser(&User{UserID: "u8", Username: "user8", Email: "u8@t.com"})) 333 s := newTestSession("sess-share", "user8", "mcp_scan") 334 require.NoError(t, ts.CreateSession(s)) 335 336 require.NoError(t, ts.SetShare("sess-share", true)) 337 got, err := ts.GetSession("sess-share") 338 require.NoError(t, err) 339 assert.True(t, got.Share) 340 } 341 342 func TestTaskStore_DeleteSessionWithMessages(t *testing.T) { 343 ts, _, cleanup := newTestDB(t) 344 defer cleanup() 345 346 require.NoError(t, ts.CreateUser(&User{UserID: "u9", Username: "user9", Email: "u9@t.com"})) 347 s := newTestSession("sess-del", "user9", "ai_infra_scan") 348 require.NoError(t, ts.CreateSession(s)) 349 350 // Add a message 351 require.NoError(t, ts.StoreEvent("msg-1", "sess-del", "liveStatus", map[string]string{"text": "hello"}, 1000)) 352 353 require.NoError(t, ts.DeleteSessionWithMessages("sess-del")) 354 _, err := ts.GetSession("sess-del") 355 assert.Error(t, err) 356 357 msgs, err := ts.GetSessionMessages("sess-del") 358 require.NoError(t, err) 359 assert.Empty(t, msgs) 360 } 361 362 func TestTaskStore_ResetRunningTasks(t *testing.T) { 363 ts, _, cleanup := newTestDB(t) 364 defer cleanup() 365 366 require.NoError(t, ts.CreateUser(&User{UserID: "u10", Username: "user10", Email: "u10@t.com"})) 367 s := newTestSession("sess-run", "user10", "mcp_scan") 368 s.Status = "doing" 369 require.NoError(t, ts.CreateSession(s)) 370 371 require.NoError(t, ts.ResetRunningTasks()) 372 373 got, err := ts.GetSession("sess-run") 374 require.NoError(t, err) 375 assert.Equal(t, "error", got.Status) 376 } 377 378 func TestTaskStore_SearchUserSessionsSimple(t *testing.T) { 379 ts, _, cleanup := newTestDB(t) 380 defer cleanup() 381 382 require.NoError(t, ts.CreateUser(&User{UserID: "u11", Username: "srchuser", Email: "srch@t.com"})) 383 s1 := newTestSession("srch-1", "srchuser", "mcp_scan") 384 s1.Title = "important scan" 385 require.NoError(t, ts.CreateSession(s1)) 386 387 s2 := newTestSession("srch-2", "srchuser", "ai_infra_scan") 388 s2.Title = "other task" 389 require.NoError(t, ts.CreateSession(s2)) 390 391 // Search by keyword 392 results, total, err := ts.SearchUserSessionsSimple("srchuser", SimpleSearchParams{ 393 Query: "important", Page: 1, PageSize: 10, 394 }) 395 require.NoError(t, err) 396 assert.Equal(t, int64(1), total) 397 assert.Len(t, results, 1) 398 assert.Equal(t, "srch-1", results[0].ID) 399 400 // Search with task type filter 401 results, total, err = ts.SearchUserSessionsSimple("srchuser", SimpleSearchParams{ 402 TaskType: "ai_infra_scan", Page: 1, PageSize: 10, 403 }) 404 require.NoError(t, err) 405 assert.Equal(t, int64(1), total) 406 assert.Equal(t, "ai_infra_scan", results[0].TaskType) 407 } 408 409 // --------------------------------------------------------------------------- 410 // TaskStore – Messages (AddMessage / GetMessages) 411 // --------------------------------------------------------------------------- 412 413 func TestTaskStore_AddAndGetMessages(t *testing.T) { 414 ts, _, cleanup := newTestDB(t) 415 defer cleanup() 416 417 require.NoError(t, ts.CreateUser(&User{UserID: "u12", Username: "msguser", Email: "msg@t.com"})) 418 require.NoError(t, ts.CreateSession(newTestSession("msg-sess", "msguser", "mcp_scan"))) 419 420 require.NoError(t, ts.StoreEvent("ev-1", "msg-sess", "liveStatus", map[string]string{"text": "starting"}, 1000)) 421 require.NoError(t, ts.StoreEvent("ev-2", "msg-sess", "actionLog", map[string]string{"actionLog": "step 1"}, 2000)) 422 require.NoError(t, ts.StoreEvent("ev-3", "msg-sess", "resultUpdate", map[string]interface{}{"result": 42}, 3000)) 423 424 msgs, err := ts.GetSessionMessages("msg-sess") 425 require.NoError(t, err) 426 assert.Len(t, msgs, 3) 427 428 // Should be ordered by timestamp ASC 429 assert.Equal(t, "ev-1", msgs[0].ID) 430 assert.Equal(t, "ev-2", msgs[1].ID) 431 assert.Equal(t, "ev-3", msgs[2].ID) 432 } 433 434 func TestTaskStore_GetSessionEventsByType(t *testing.T) { 435 ts, _, cleanup := newTestDB(t) 436 defer cleanup() 437 438 require.NoError(t, ts.CreateUser(&User{UserID: "u13", Username: "typeuser", Email: "type@t.com"})) 439 require.NoError(t, ts.CreateSession(newTestSession("type-sess", "typeuser", "mcp_scan"))) 440 441 require.NoError(t, ts.StoreEvent("t1", "type-sess", "liveStatus", map[string]string{"text": "a"}, 100)) 442 require.NoError(t, ts.StoreEvent("t2", "type-sess", "actionLog", map[string]string{"actionLog": "b"}, 200)) 443 require.NoError(t, ts.StoreEvent("t3", "type-sess", "liveStatus", map[string]string{"text": "c"}, 300)) 444 445 lsEvents, err := ts.GetSessionEventsByType("type-sess", "liveStatus") 446 require.NoError(t, err) 447 assert.Len(t, lsEvents, 2) 448 449 alEvents, err := ts.GetSessionEventsByType("type-sess", "actionLog") 450 require.NoError(t, err) 451 assert.Len(t, alEvents, 1) 452 } 453 454 func TestTaskStore_CreateTaskMessage(t *testing.T) { 455 ts, _, cleanup := newTestDB(t) 456 defer cleanup() 457 458 require.NoError(t, ts.CreateUser(&User{UserID: "u14", Username: "tmuser", Email: "tm@t.com"})) 459 require.NoError(t, ts.CreateSession(newTestSession("tm-sess", "tmuser", "mcp_scan"))) 460 461 msg := &TaskMessage{ 462 ID: "manual-msg", 463 SessionID: "tm-sess", 464 Type: "planUpdate", 465 EventData: []byte(`{"tasks":[]}`), 466 Timestamp: 999, 467 } 468 require.NoError(t, ts.CreateTaskMessage(msg)) 469 assert.Greater(t, msg.CreatedAt, int64(0)) 470 471 msgs, err := ts.GetSessionEvents("tm-sess") 472 require.NoError(t, err) 473 assert.Len(t, msgs, 1) 474 assert.Equal(t, "manual-msg", msgs[0].ID) 475 } 476 477 // --------------------------------------------------------------------------- 478 // ModelStore – Init, basic CRUD 479 // --------------------------------------------------------------------------- 480 481 func TestModelStore_Init(t *testing.T) { 482 _, ms, cleanup := newTestDB(t) 483 defer cleanup() 484 // Calling Init twice should be idempotent 485 assert.NoError(t, ms.Init()) 486 } 487 488 func TestModelStore_CreateAndGetModel(t *testing.T) { 489 ts, ms, cleanup := newTestDB(t) 490 defer cleanup() 491 492 require.NoError(t, ts.CreateUser(&User{UserID: "mu1", Username: "muser1", Email: "muser1@t.com"})) 493 494 m := &Model{ 495 ModelID: "model-001", 496 Username: "muser1", 497 ModelName: "gpt-4", 498 Token: "sk-test", 499 BaseURL: "https://api.openai.com/v1", 500 } 501 require.NoError(t, ms.CreateModel(m)) 502 assert.Greater(t, m.CreatedAt, int64(0)) 503 assert.Greater(t, m.UpdatedAt, int64(0)) 504 505 got, err := ms.GetModel("model-001") 506 require.NoError(t, err) 507 assert.Equal(t, "gpt-4", got.ModelName) 508 assert.Equal(t, "sk-test", got.Token) 509 } 510 511 func TestModelStore_GetModel_NotFound(t *testing.T) { 512 _, ms, cleanup := newTestDB(t) 513 defer cleanup() 514 515 _, err := ms.GetModel("nonexistent-model") 516 assert.Error(t, err) 517 } 518 519 func TestModelStore_GetModelByUser(t *testing.T) { 520 ts, ms, cleanup := newTestDB(t) 521 defer cleanup() 522 523 require.NoError(t, ts.CreateUser(&User{UserID: "mu2", Username: "muser2", Email: "muser2@t.com"})) 524 m := &Model{ 525 ModelID: "model-002", 526 Username: "muser2", 527 ModelName: "gpt-3.5", 528 Token: "sk-test2", 529 BaseURL: "https://api.openai.com/v1", 530 } 531 require.NoError(t, ms.CreateModel(m)) 532 533 got, err := ms.GetModelByUser("model-002", "muser2") 534 require.NoError(t, err) 535 assert.Equal(t, "gpt-3.5", got.ModelName) 536 537 _, err = ms.GetModelByUser("model-002", "wronguser") 538 assert.Error(t, err) 539 } 540 541 func TestModelStore_GetAllModels(t *testing.T) { 542 ts, ms, cleanup := newTestDB(t) 543 defer cleanup() 544 545 require.NoError(t, ts.CreateUser(&User{UserID: "mu3", Username: "muser3", Email: "muser3@t.com"})) 546 for i, id := range []string{"m1", "m2", "m3"} { 547 require.NoError(t, ms.CreateModel(&Model{ 548 ModelID: id, 549 Username: "muser3", 550 ModelName: "model-" + string(rune('A'+i)), 551 Token: "tok", 552 BaseURL: "https://example.com", 553 })) 554 } 555 556 all, err := ms.GetAllModels() 557 require.NoError(t, err) 558 assert.Len(t, all, 3) 559 } 560 561 func TestModelStore_UpdateModel(t *testing.T) { 562 ts, ms, cleanup := newTestDB(t) 563 defer cleanup() 564 565 require.NoError(t, ts.CreateUser(&User{UserID: "mu4", Username: "muser4", Email: "muser4@t.com"})) 566 require.NoError(t, ms.CreateModel(&Model{ 567 ModelID: "model-upd", 568 Username: "muser4", 569 ModelName: "old-model", 570 Token: "tok", 571 BaseURL: "https://example.com", 572 })) 573 574 require.NoError(t, ms.UpdateModel("model-upd", "muser4", map[string]interface{}{"model_name": "new-model"})) 575 got, err := ms.GetModel("model-upd") 576 require.NoError(t, err) 577 assert.Equal(t, "new-model", got.ModelName) 578 } 579 580 func TestModelStore_DeleteModel(t *testing.T) { 581 ts, ms, cleanup := newTestDB(t) 582 defer cleanup() 583 584 require.NoError(t, ts.CreateUser(&User{UserID: "mu5", Username: "muser5", Email: "muser5@t.com"})) 585 require.NoError(t, ms.CreateModel(&Model{ 586 ModelID: "model-del", 587 Username: "muser5", 588 ModelName: "to-delete", 589 Token: "tok", 590 BaseURL: "https://example.com", 591 })) 592 593 require.NoError(t, ms.DeleteModel("model-del", "muser5")) 594 595 exists, err := ms.CheckModelExists("model-del") 596 require.NoError(t, err) 597 assert.False(t, exists) 598 } 599 600 func TestModelStore_BatchDeleteModels(t *testing.T) { 601 ts, ms, cleanup := newTestDB(t) 602 defer cleanup() 603 604 require.NoError(t, ts.CreateUser(&User{UserID: "mu6", Username: "muser6", Email: "muser6@t.com"})) 605 ids := []string{"bm1", "bm2", "bm3"} 606 for _, id := range ids { 607 require.NoError(t, ms.CreateModel(&Model{ 608 ModelID: id, 609 Username: "muser6", 610 ModelName: id, 611 Token: "tok", 612 BaseURL: "https://example.com", 613 })) 614 } 615 616 n, err := ms.BatchDeleteModels(ids, "muser6") 617 require.NoError(t, err) 618 assert.Equal(t, int64(3), n) 619 } 620 621 func TestModelStore_CheckModelExists(t *testing.T) { 622 ts, ms, cleanup := newTestDB(t) 623 defer cleanup() 624 625 require.NoError(t, ts.CreateUser(&User{UserID: "mu7", Username: "muser7", Email: "muser7@t.com"})) 626 require.NoError(t, ms.CreateModel(&Model{ 627 ModelID: "model-chk", 628 Username: "muser7", 629 ModelName: "chk", 630 Token: "tok", 631 BaseURL: "https://example.com", 632 })) 633 634 exists, err := ms.CheckModelExists("model-chk") 635 require.NoError(t, err) 636 assert.True(t, exists) 637 638 exists, err = ms.CheckModelExists("nope") 639 require.NoError(t, err) 640 assert.False(t, exists) 641 } 642 643 func TestModelStore_CheckModelExistsByUser(t *testing.T) { 644 ts, ms, cleanup := newTestDB(t) 645 defer cleanup() 646 647 require.NoError(t, ts.CreateUser(&User{UserID: "mu8", Username: "muser8", Email: "muser8@t.com"})) 648 require.NoError(t, ms.CreateModel(&Model{ 649 ModelID: "model-u", 650 Username: "muser8", 651 ModelName: "u", 652 Token: "tok", 653 BaseURL: "https://example.com", 654 })) 655 656 exists, err := ms.CheckModelExistsByUser("model-u", "muser8") 657 require.NoError(t, err) 658 assert.True(t, exists) 659 660 exists, err = ms.CheckModelExistsByUser("model-u", "other") 661 require.NoError(t, err) 662 assert.False(t, exists) 663 } 664 665 func TestModelStore_GetUserModels(t *testing.T) { 666 ts, ms, cleanup := newTestDB(t) 667 defer cleanup() 668 669 require.NoError(t, ts.CreateUser(&User{UserID: "mu9", Username: "muser9", Email: "muser9@t.com"})) 670 for _, id := range []string{"um1", "um2"} { 671 require.NoError(t, ms.CreateModel(&Model{ 672 ModelID: id, 673 Username: "muser9", 674 ModelName: id, 675 Token: "tok", 676 BaseURL: "https://example.com", 677 })) 678 } 679 680 models, err := ms.GetUserModels("muser9") 681 require.NoError(t, err) 682 // At least the two created (yaml models may also appear if file exists) 683 assert.GreaterOrEqual(t, len(models), 2) 684 }