115 lines
4.2 KiB
Swift
115 lines
4.2 KiB
Swift
@testable import MusicAssistantCore
|
|
import XCTest
|
|
|
|
final class ApplicationRuleInjectionTests: XCTestCase {
|
|
func testInjectingServiceAddsPrivateRulesToEveryAIRequestType() async throws {
|
|
let baseService = RecordingAIService()
|
|
let service = ApplicationRuleInjectingAIService(
|
|
baseService: baseService,
|
|
ruleProvider: StaticApplicationRuleProvider(
|
|
rules: PrivateApplicationRules(content: "Use private product rules.")
|
|
)
|
|
)
|
|
let project = SongProject(title: "Rule Test", idea: "Test rules")
|
|
|
|
_ = try await service.generateSongProject(
|
|
from: SongProjectGenerationRequest(
|
|
context: AIRequestContext(userInstruction: "Generate.")
|
|
)
|
|
)
|
|
_ = try await service.discussSongProject(
|
|
from: SongProjectDiscussionRequest(
|
|
context: AIRequestContext(userInstruction: "Discuss."),
|
|
project: project
|
|
)
|
|
)
|
|
_ = try await service.reviseLyrics(
|
|
from: LyricsRevisionRequest(
|
|
context: AIRequestContext(userInstruction: "Improve."),
|
|
project: project,
|
|
sourceLyrics: Lyrics(text: "Draft"),
|
|
mode: .improve
|
|
)
|
|
)
|
|
_ = try await service.proposeProjectUpdate(
|
|
from: SongProjectUpdateRequest(
|
|
context: AIRequestContext(userInstruction: "Update."),
|
|
project: project,
|
|
allowedScopes: [.lyrics]
|
|
)
|
|
)
|
|
|
|
let recordedRuleContents = await baseService.recordedRuleContents()
|
|
XCTAssertEqual(
|
|
recordedRuleContents,
|
|
[
|
|
"Use private product rules.",
|
|
"Use private product rules.",
|
|
"Use private product rules.",
|
|
"Use private product rules."
|
|
]
|
|
)
|
|
}
|
|
|
|
func testInjectingServiceDoesNotAddPrivateRulesToVisibleConversation() async throws {
|
|
let baseService = RecordingAIService()
|
|
let service = ApplicationRuleInjectingAIService(
|
|
baseService: baseService,
|
|
ruleProvider: StaticApplicationRuleProvider(
|
|
rules: PrivateApplicationRules(content: "Hidden owner rules")
|
|
)
|
|
)
|
|
let visibleMessage = AIConversationMessage(role: .user, content: "Visible user text.")
|
|
|
|
_ = try await service.generateSongProject(
|
|
from: SongProjectGenerationRequest(
|
|
context: AIRequestContext(
|
|
userInstruction: "Generate.",
|
|
conversation: [visibleMessage]
|
|
)
|
|
)
|
|
)
|
|
|
|
let recordedConversation = await baseService.recordedConversations().first
|
|
XCTAssertEqual(recordedConversation, [visibleMessage])
|
|
}
|
|
|
|
func testBlankPrivateRulesAreIgnored() {
|
|
XCTAssertNil(PrivateApplicationRules(content: " \n "))
|
|
}
|
|
}
|
|
|
|
private actor RecordingAIService: AIService {
|
|
private var contexts: [AIRequestContext] = []
|
|
|
|
func recordedRuleContents() -> [String?] {
|
|
contexts.map { $0.privateApplicationRules?.content }
|
|
}
|
|
|
|
func recordedConversations() -> [[AIConversationMessage]] {
|
|
contexts.map(\.conversation)
|
|
}
|
|
|
|
func generateSongProject(from request: SongProjectGenerationRequest) async throws -> SongProjectGenerationResult {
|
|
contexts.append(request.context)
|
|
return SongProjectGenerationResult(
|
|
project: SongProject(title: "Generated", idea: request.context.userInstruction)
|
|
)
|
|
}
|
|
|
|
func discussSongProject(from request: SongProjectDiscussionRequest) async throws -> SongProjectDiscussionResult {
|
|
contexts.append(request.context)
|
|
return SongProjectDiscussionResult(questions: ["Which direction should we take?"])
|
|
}
|
|
|
|
func reviseLyrics(from request: LyricsRevisionRequest) async throws -> LyricsRevisionResult {
|
|
contexts.append(request.context)
|
|
return LyricsRevisionResult(lyrics: request.sourceLyrics)
|
|
}
|
|
|
|
func proposeProjectUpdate(from request: SongProjectUpdateRequest) async throws -> SongProjectUpdateResult {
|
|
contexts.append(request.context)
|
|
return SongProjectUpdateResult(project: request.project)
|
|
}
|
|
}
|