diff --git a/features/image-assistant/backend-api/src/modules/classification/classification.logic.spec.ts b/features/image-assistant/backend-api/src/modules/classification/classification.logic.spec.ts new file mode 100644 index 000000000..31d08a65c --- /dev/null +++ b/features/image-assistant/backend-api/src/modules/classification/classification.logic.spec.ts @@ -0,0 +1,155 @@ +import { describe, it, expect } from 'vitest'; + +import { determineCategory, SCREENSHOT_SIGLIP_PROMPTS } from './classification.logic'; +import { PhotoCategory } from './classification.types'; + +import type { ModeratorResult, SiglipResult } from './classification.types'; + +describe('SCREENSHOT_SIGLIP_PROMPTS', () => { + it('contains all expected screenshot categories as keys', () => { + const expectedKeys = ['receipt', 'conversation', 'shopping', 'meme', 'event', 'reservation', 'hottie', 'other']; + for (const key of expectedKeys) { + expect(SCREENSHOT_SIGLIP_PROMPTS).toHaveProperty(key); + expect(Array.isArray(SCREENSHOT_SIGLIP_PROMPTS[key])).toBe(true); + expect(SCREENSHOT_SIGLIP_PROMPTS[key].length).toBeGreaterThan(0); + } + }); +}); + +describe('determineCategory — screenshot path', () => { + const screenshotPhoto = { isScreenshot: true, isSelfie: false, storageKey: 'some/key.jpg' }; + + it('returns SCREENSHOT_OTHER when no siglipResult provided', () => { + expect(determineCategory(screenshotPhoto)).toBe(PhotoCategory.SCREENSHOT_OTHER); + }); + + it('returns SCREENSHOT_OTHER when siglipResult has empty scores', () => { + const siglip: SiglipResult = { scores: {} }; + expect(determineCategory(screenshotPhoto, undefined, siglip)).toBe(PhotoCategory.SCREENSHOT_OTHER); + }); + + it('picks the key with the highest score', () => { + const siglip: SiglipResult = { + scores: { receipt: 0.1, conversation: 0.9, shopping: 0.3 }, + }; + expect(determineCategory(screenshotPhoto, undefined, siglip)).toBe(PhotoCategory.SCREENSHOT_CONVERSATION); + }); + + it('maps receipt → SCREENSHOT_RECEIPT', () => { + const siglip: SiglipResult = { scores: { receipt: 0.95, other: 0.1 } }; + expect(determineCategory(screenshotPhoto, undefined, siglip)).toBe(PhotoCategory.SCREENSHOT_RECEIPT); + }); + + it('maps shopping → SCREENSHOT_SHOPPING', () => { + const siglip: SiglipResult = { scores: { shopping: 0.8 } }; + expect(determineCategory(screenshotPhoto, undefined, siglip)).toBe(PhotoCategory.SCREENSHOT_SHOPPING); + }); + + it('maps meme → SCREENSHOT_MEME', () => { + const siglip: SiglipResult = { scores: { meme: 0.7 } }; + expect(determineCategory(screenshotPhoto, undefined, siglip)).toBe(PhotoCategory.SCREENSHOT_MEME); + }); + + it('maps event → SCREENSHOT_EVENT', () => { + const siglip: SiglipResult = { scores: { event: 0.6 } }; + expect(determineCategory(screenshotPhoto, undefined, siglip)).toBe(PhotoCategory.SCREENSHOT_EVENT); + }); + + it('maps reservation → SCREENSHOT_RESERVATION', () => { + const siglip: SiglipResult = { scores: { reservation: 0.75 } }; + expect(determineCategory(screenshotPhoto, undefined, siglip)).toBe(PhotoCategory.SCREENSHOT_RESERVATION); + }); + + it('maps hottie → SCREENSHOT_HOTTIE', () => { + const siglip: SiglipResult = { scores: { hottie: 0.85 } }; + expect(determineCategory(screenshotPhoto, undefined, siglip)).toBe(PhotoCategory.SCREENSHOT_HOTTIE); + }); + + it('maps other → SCREENSHOT_OTHER', () => { + const siglip: SiglipResult = { scores: { other: 0.6 } }; + expect(determineCategory(screenshotPhoto, undefined, siglip)).toBe(PhotoCategory.SCREENSHOT_OTHER); + }); + + it('returns SCREENSHOT_OTHER for an unknown top key', () => { + const siglip: SiglipResult = { scores: { unknown_category: 0.99 } }; + expect(determineCategory(screenshotPhoto, undefined, siglip)).toBe(PhotoCategory.SCREENSHOT_OTHER); + }); + + it('ignores moderatorResult when photo is a screenshot', () => { + const moderator: ModeratorResult = { nsfw_score: 0.9, nsfw_category: 'explicit', face_count: 2 }; + const siglip: SiglipResult = { scores: { meme: 0.9 } }; + expect(determineCategory(screenshotPhoto, moderator, siglip)).toBe(PhotoCategory.SCREENSHOT_MEME); + }); +}); + +describe('determineCategory — no storage key', () => { + it('returns UNCLASSIFIED when storageKey is null and not a screenshot', () => { + expect(determineCategory({ isScreenshot: false, isSelfie: false, storageKey: null })).toBe( + PhotoCategory.UNCLASSIFIED, + ); + }); + + it('returns UNCLASSIFIED when storageKey is undefined and not a screenshot', () => { + expect(determineCategory({ isScreenshot: false, isSelfie: false })).toBe(PhotoCategory.UNCLASSIFIED); + }); +}); + +describe('determineCategory — moderator result path', () => { + const photo = { isScreenshot: false, isSelfie: false, storageKey: 'photos/original.jpg' }; + const selfiePhoto = { isScreenshot: false, isSelfie: true, storageKey: 'photos/selfie.jpg' }; + + it('returns UNCLASSIFIED when no moderatorResult is provided', () => { + expect(determineCategory(photo)).toBe(PhotoCategory.UNCLASSIFIED); + }); + + it('returns SELF_EXPLICIT for explicit + selfie', () => { + const moderator: ModeratorResult = { nsfw_score: 0.95, nsfw_category: 'explicit', face_count: 1 }; + expect(determineCategory(selfiePhoto, moderator)).toBe(PhotoCategory.SELF_EXPLICIT); + }); + + it('returns SELF_EXPLICIT for explicit + face_count <= 1 even without isSelfie flag', () => { + const moderator: ModeratorResult = { nsfw_score: 0.95, nsfw_category: 'explicit', face_count: 1 }; + expect(determineCategory(photo, moderator)).toBe(PhotoCategory.SELF_EXPLICIT); + }); + + it('returns SELF_WITH_OTHERS for explicit + face_count > 1 without isSelfie', () => { + const moderator: ModeratorResult = { nsfw_score: 0.95, nsfw_category: 'explicit', face_count: 3 }; + expect(determineCategory(photo, moderator)).toBe(PhotoCategory.SELF_WITH_OTHERS); + }); + + it('returns SELF_EXPLICIT for explicit + isSelfie even with face_count > 1', () => { + const moderator: ModeratorResult = { nsfw_score: 0.95, nsfw_category: 'explicit', face_count: 2 }; + expect(determineCategory(selfiePhoto, moderator)).toBe(PhotoCategory.SELF_EXPLICIT); + }); + + it('returns SELF_NUDE when nsfw_score > 0.4 and not explicit', () => { + const moderator: ModeratorResult = { nsfw_score: 0.6, nsfw_category: 'nude', face_count: 0 }; + expect(determineCategory(photo, moderator)).toBe(PhotoCategory.SELF_NUDE); + }); + + it('returns SELF_NUDE at exactly the 0.4 threshold boundary (above)', () => { + const moderator: ModeratorResult = { nsfw_score: 0.41, nsfw_category: 'suggestive', face_count: 0 }; + expect(determineCategory(photo, moderator)).toBe(PhotoCategory.SELF_NUDE); + }); + + it('does NOT return SELF_NUDE at exactly 0.4 (boundary is strict greater-than)', () => { + const moderator: ModeratorResult = { nsfw_score: 0.4, nsfw_category: 'safe', face_count: 0 }; + // Falls through to isSelfie/face_count checks + expect(determineCategory(photo, moderator)).toBe(PhotoCategory.UNCLASSIFIED); + }); + + it('returns SELF_CLOTHED when isSelfie and nsfw_score <= 0.4', () => { + const moderator: ModeratorResult = { nsfw_score: 0.1, nsfw_category: 'safe', face_count: 1 }; + expect(determineCategory(selfiePhoto, moderator)).toBe(PhotoCategory.SELF_CLOTHED); + }); + + it('returns FRIENDS when face_count > 0 and not selfie and nsfw_score <= 0.4', () => { + const moderator: ModeratorResult = { nsfw_score: 0.1, nsfw_category: 'safe', face_count: 2 }; + expect(determineCategory(photo, moderator)).toBe(PhotoCategory.FRIENDS); + }); + + it('returns UNCLASSIFIED when safe, no selfie, no faces', () => { + const moderator: ModeratorResult = { nsfw_score: 0.05, nsfw_category: 'safe', face_count: 0 }; + expect(determineCategory(photo, moderator)).toBe(PhotoCategory.UNCLASSIFIED); + }); +});