oai-swift/oAITests/ContextSelectionServiceDBTests.swift

//
//  ContextSelectionServiceDBTests.swift
//  oAITests
//
// SPDX-License-Identifier: PolyForm-Noncommercial-1.0.0
// Copyright (C) 2026 Rune Olsen

import Testing
import Foundation
@testable import oAI

@Suite("ContextSelectionService DB-coupled paths, against a throwaway in-memory queue")
struct ContextSelectionServiceDBTests {

    /// Persists messages, then loads them back so their `id`s match the DB rows
    /// (saveConversation assigns fresh row ids internally, distinct from the ids
    /// on the Message values passed in).
    private func persistedMessages(_ db: DatabaseService, _ messages: [Message]) throws -> (conversationId: UUID, messages: [Message]) {
        let saved = try db.saveConversation(name: "Test", messages: messages)
        let loaded = try db.loadConversation(id: saved.id)
        return (saved.id, loaded?.1 ?? [])
    }

    // MARK: - isMessageStarred

    @Test("A message with no metadata row is not starred")
    func unstarredByDefault() throws {
        let db = DatabaseService.makeInMemory()
        let service = ContextSelectionService(db: db)
        let (_, messages) = try persistedMessages(db, [Message(role: .user, content: "hi")])

        #expect(service.isMessageStarred(messages[0]) == false)
    }

    @Test("A message starred via setMessageStarred reports starred")
    func starredAfterSetting() throws {
        let db = DatabaseService.makeInMemory()
        let service = ContextSelectionService(db: db)
        let (_, messages) = try persistedMessages(db, [Message(role: .user, content: "hi")])

        try db.setMessageStarred(messageId: messages[0].id, starred: true)
        #expect(service.isMessageStarred(messages[0]) == true)
    }

    // MARK: - getSummariesForExcludedRange

    @Test("No summaries recorded returns an empty array")
    func noSummariesReturnsEmpty() throws {
        let db = DatabaseService.makeInMemory()
        let service = ContextSelectionService(db: db)
        let (conversationId, _) = try persistedMessages(db, [Message(role: .user, content: "hi")])

        let summaries = service.getSummariesForExcludedRange(conversationId: conversationId, totalMessages: 30, selectedCount: 10)
        #expect(summaries.isEmpty)
    }

    @Test("A summary covering only excluded messages is included")
    func summaryForExcludedRangeIncluded() throws {
        let db = DatabaseService.makeInMemory()
        let service = ContextSelectionService(db: db)
        let (conversationId, _) = try persistedMessages(db, [Message(role: .user, content: "hi")])

        // 30 total, 10 selected -> messages 0..<20 are excluded
        try db.saveConversationSummary(conversationId: conversationId, startIndex: 0, endIndex: 15, summary: "early stuff", model: nil, tokenCount: nil)

        let summaries = service.getSummariesForExcludedRange(conversationId: conversationId, totalMessages: 30, selectedCount: 10)
        #expect(summaries == ["early stuff"])
    }

    @Test("A summary covering messages that were kept is excluded")
    func summaryForKeptRangeExcluded() throws {
        let db = DatabaseService.makeInMemory()
        let service = ContextSelectionService(db: db)
        let (conversationId, _) = try persistedMessages(db, [Message(role: .user, content: "hi")])

        // 30 total, 10 selected -> messages 0..<20 are excluded; this summary ends at 25, inside the kept range
        try db.saveConversationSummary(conversationId: conversationId, startIndex: 15, endIndex: 25, summary: "later stuff", model: nil, tokenCount: nil)

        let summaries = service.getSummariesForExcludedRange(conversationId: conversationId, totalMessages: 30, selectedCount: 10)
        #expect(summaries.isEmpty)
    }

    // MARK: - smartSelection (end to end through the DB seam)

    @Test("A starred older message is pulled in alongside the recent window")
    func starredMessagePulledIntoSelection() throws {
        let db = DatabaseService.makeInMemory()
        let service = ContextSelectionService(db: db)

        var all: [Message] = []
        let base = Date()
        for i in 0..<15 {
            all.append(Message(role: i % 2 == 0 ? .user : .assistant, content: "msg \(i)", timestamp: base.addingTimeInterval(Double(i))))
        }
        let (_, persisted) = try persistedMessages(db, all)

        // Star an old message that's outside the last-10 window (index 2)
        try db.setMessageStarred(messageId: persisted[2].id, starred: true)

        let window = service.smartSelection(allMessages: persisted, maxTokens: 100_000, conversationId: nil)
        #expect(window.messages.contains { $0.id == persisted[2].id })
    }

    @Test("An unstarred, unimportant old message outside the recent window is excluded")
    func unstarredOldMessageExcluded() throws {
        let db = DatabaseService.makeInMemory()
        let service = ContextSelectionService(db: db)

        var all: [Message] = []
        let base = Date()
        for i in 0..<15 {
            all.append(Message(role: i % 2 == 0 ? .user : .assistant, content: "msg \(i)", timestamp: base.addingTimeInterval(Double(i))))
        }
        let (_, persisted) = try persistedMessages(db, all)

        let window = service.smartSelection(allMessages: persisted, maxTokens: 100_000, conversationId: nil)
        #expect(window.messages.contains { $0.id == persisted[0].id } == false)
        #expect(window.excludedCount == 5)
    }
}