|
1 | 1 | package repository |
2 | 2 |
|
3 | 3 | import ( |
| 4 | + "encoding/json" |
4 | 5 | "sync/atomic" |
5 | 6 | "testing" |
6 | 7 | "time" |
@@ -579,3 +580,100 @@ func TestFaceRepository_ListPersonIDsSharingPhotos_NoNPlusOne(t *testing.T) { |
579 | 580 | assert.Less(t, atomic.LoadInt32(qcount), int32(600), "no per-candidate N+1; chunked and bounded") |
580 | 581 | assert.Greater(t, atomic.LoadInt32(qcount), int32(0)) |
581 | 582 | } |
| 583 | + |
| 584 | +func TestFaceRepository_ListAssignedPersonIDsPaged(t *testing.T) { |
| 585 | + db := setupTestDB(t) |
| 586 | + defer teardownTestDB(db) |
| 587 | + faceRepo := NewFaceRepository(db) |
| 588 | + personRepo := NewPersonRepository(db) |
| 589 | + |
| 590 | + // Create 5 persons with faces assigned. |
| 591 | + personIDs := make([]uint, 5) |
| 592 | + for i := 0; i < 5; i++ { |
| 593 | + p := &model.Person{Category: model.PersonCategoryFriend} |
| 594 | + require.NoError(t, personRepo.Create(p)) |
| 595 | + personIDs[i] = p.ID |
| 596 | + require.NoError(t, faceRepo.Create(&model.Face{ |
| 597 | + PhotoID: uint(i + 1), |
| 598 | + PersonID: &p.ID, |
| 599 | + BBoxX: 0.1, BBoxY: 0.1, BBoxWidth: 0.2, BBoxHeight: 0.2, |
| 600 | + Confidence: 0.9, QualityScore: 0.8, |
| 601 | + })) |
| 602 | + } |
| 603 | + // Also create an unassigned face — should not appear. |
| 604 | + require.NoError(t, faceRepo.Create(&model.Face{PhotoID: 99, BBoxX: 0.1, BBoxY: 0.1, BBoxWidth: 0.2, BBoxHeight: 0.2})) |
| 605 | + |
| 606 | + // Page 1: offset=0, limit=3 → first 3 person IDs (ascending). |
| 607 | + page1, err := faceRepo.ListAssignedPersonIDsPaged(0, 3) |
| 608 | + require.NoError(t, err) |
| 609 | + assert.Len(t, page1, 3) |
| 610 | + assert.Equal(t, personIDs[0], page1[0]) |
| 611 | + assert.Equal(t, personIDs[1], page1[1]) |
| 612 | + assert.Equal(t, personIDs[2], page1[2]) |
| 613 | + |
| 614 | + // Page 2: offset=3, limit=3 → last 2 person IDs. |
| 615 | + page2, err := faceRepo.ListAssignedPersonIDsPaged(3, 3) |
| 616 | + require.NoError(t, err) |
| 617 | + assert.Len(t, page2, 2) |
| 618 | + assert.Equal(t, personIDs[3], page2[0]) |
| 619 | + assert.Equal(t, personIDs[4], page2[1]) |
| 620 | + |
| 621 | + // Page 3: offset=5, limit=3 → empty. |
| 622 | + page3, err := faceRepo.ListAssignedPersonIDsPaged(5, 3) |
| 623 | + require.NoError(t, err) |
| 624 | + assert.Empty(t, page3) |
| 625 | + |
| 626 | + // limit=0 → empty (no error). |
| 627 | + empty, err := faceRepo.ListAssignedPersonIDsPaged(0, 0) |
| 628 | + require.NoError(t, err) |
| 629 | + assert.Empty(t, empty) |
| 630 | +} |
| 631 | + |
| 632 | +func TestFaceRepository_ListPrototypeEmbeddings_Batched(t *testing.T) { |
| 633 | + db := setupTestDB(t) |
| 634 | + defer teardownTestDB(db) |
| 635 | + faceRepo := NewFaceRepository(db) |
| 636 | + personRepo := NewPersonRepository(db) |
| 637 | + |
| 638 | + // Create 2 persons, each with 3 faces (different quality). |
| 639 | + emb := encodeFloat32(t, []float32{1.0, 0.0, 0.0}) |
| 640 | + for pid := 1; pid <= 2; pid++ { |
| 641 | + p := &model.Person{Category: model.PersonCategoryFriend} |
| 642 | + require.NoError(t, personRepo.Create(p)) |
| 643 | + for f := 0; f < 3; f++ { |
| 644 | + require.NoError(t, faceRepo.Create(&model.Face{ |
| 645 | + PhotoID: uint(pid*10 + f), |
| 646 | + PersonID: &p.ID, |
| 647 | + BBoxX: 0.1, BBoxY: 0.1, BBoxWidth: 0.2, BBoxHeight: 0.2, |
| 648 | + Confidence: 0.9, |
| 649 | + QualityScore: float64(3 - f), // 3, 2, 1 |
| 650 | + Embedding: emb, |
| 651 | + })) |
| 652 | + } |
| 653 | + } |
| 654 | + |
| 655 | + personIDs, err := faceRepo.ListAssignedPersonIDsPaged(0, 10) |
| 656 | + require.NoError(t, err) |
| 657 | + assert.Len(t, personIDs, 2) |
| 658 | + |
| 659 | + // Request top 2 per person → 4 faces total. |
| 660 | + faces, err := faceRepo.ListPrototypeEmbeddings(personIDs, 2) |
| 661 | + require.NoError(t, err) |
| 662 | + assert.Len(t, faces, 4) // 2 persons × 2 per person |
| 663 | + |
| 664 | + // Verify each person got 2 faces (highest quality first). |
| 665 | + byPerson := make(map[uint]int) |
| 666 | + for _, f := range faces { |
| 667 | + byPerson[*f.PersonID]++ |
| 668 | + } |
| 669 | + for _, count := range byPerson { |
| 670 | + assert.Equal(t, 2, count) |
| 671 | + } |
| 672 | +} |
| 673 | + |
| 674 | +func encodeFloat32(t *testing.T, vals []float32) []byte { |
| 675 | + t.Helper() |
| 676 | + payload, err := json.Marshal(vals) |
| 677 | + require.NoError(t, err) |
| 678 | + return payload |
| 679 | +} |
0 commit comments