932 lines
31 KiB
Rust
932 lines
31 KiB
Rust
use chrono::Utc;
|
|
use galaxy_graphql::scalars::time::ServerTimestamp;
|
|
use galaxyui::{App, SingletonEntity};
|
|
use settings_value::SettingsValue;
|
|
|
|
use super::*;
|
|
use crate::ai::request_usage_model::{RequestLimitInfo, RequestLimitRefreshDuration};
|
|
use crate::auth::AuthStateProvider;
|
|
use crate::test_util::settings::initialize_settings_for_tests;
|
|
use crate::workspaces::user_workspaces::UserWorkspaces;
|
|
|
|
fn create_test_request_limit_info(
|
|
limit: usize,
|
|
used: usize,
|
|
next_refresh: DateTime<Utc>,
|
|
is_unlimited: bool,
|
|
refresh_duration: RequestLimitRefreshDuration,
|
|
) -> RequestLimitInfo {
|
|
RequestLimitInfo {
|
|
limit,
|
|
num_requests_used_since_refresh: used,
|
|
next_refresh_time: ServerTimestamp::new(next_refresh),
|
|
is_unlimited,
|
|
request_limit_refresh_duration: refresh_duration,
|
|
is_unlimited_voice: false,
|
|
voice_request_limit: 0,
|
|
voice_requests_used_since_last_refresh: 0,
|
|
is_unlimited_codebase_indices: false,
|
|
max_codebase_indices: 0,
|
|
max_files_per_repo: 5000,
|
|
embedding_generation_batch_size: 100,
|
|
}
|
|
}
|
|
|
|
fn add_ai_enablement_dependencies_for_test(app: &mut App) {
|
|
app.add_singleton_model(|_| AuthStateProvider::new_for_test());
|
|
app.add_singleton_model(UserWorkspaces::default_mock);
|
|
}
|
|
|
|
// FocusedTerminalInfo Tests
|
|
|
|
#[test]
|
|
fn test_update_both_values_changed() {
|
|
App::test((), |mut app| async move {
|
|
// Create FocusedTerminalInfo with default values (false, false)
|
|
let model_handle = app.add_model(|_| FocusedTerminalInfo::default());
|
|
|
|
// Setup event tracking
|
|
let (sender, receiver) = async_channel::unbounded();
|
|
app.update(|ctx| {
|
|
let sender = sender.clone();
|
|
ctx.subscribe_to_model(
|
|
&model_handle,
|
|
move |_, event: &FocusedTerminalInfoEvent, _| match event {
|
|
FocusedTerminalInfoEvent::TerminalInfoUpdated => {
|
|
let _ = sender.try_send(());
|
|
}
|
|
},
|
|
);
|
|
});
|
|
|
|
// Update both values to (true, false)
|
|
model_handle.update(&mut app, |model, ctx| {
|
|
model.update(true, false, ctx);
|
|
});
|
|
|
|
// Verify model state
|
|
model_handle.read(&app, |model, _| {
|
|
assert!(model.contains_any_remote_blocks());
|
|
assert!(!model.contains_any_restored_remote_blocks());
|
|
});
|
|
|
|
// Verify event was emitted exactly once
|
|
let mut count = 0;
|
|
while receiver.try_recv().is_ok() {
|
|
count += 1;
|
|
}
|
|
assert_eq!(count, 1);
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_update_additional_value_changed() {
|
|
App::test((), |mut app| async move {
|
|
// Create FocusedTerminalInfo with default values (false, false)
|
|
let model_handle = app.add_model(|_| FocusedTerminalInfo::default());
|
|
|
|
// Setup event tracking
|
|
let (sender, receiver) = async_channel::unbounded();
|
|
app.update(|ctx| {
|
|
let sender = sender.clone();
|
|
ctx.subscribe_to_model(
|
|
&model_handle,
|
|
move |_, event: &FocusedTerminalInfoEvent, _| match event {
|
|
FocusedTerminalInfoEvent::TerminalInfoUpdated => {
|
|
let _ = sender.try_send(());
|
|
}
|
|
},
|
|
);
|
|
});
|
|
|
|
// First update to (true, false)
|
|
model_handle.update(&mut app, |model, ctx| {
|
|
model.update(true, false, ctx);
|
|
});
|
|
|
|
// Clear events by draining the channel
|
|
while receiver.try_recv().is_ok() {}
|
|
|
|
// Now update to (true, true) - only changing restored blocks
|
|
model_handle.update(&mut app, |model, ctx| {
|
|
model.update(true, true, ctx);
|
|
});
|
|
|
|
// Verify model state
|
|
model_handle.read(&app, |model, _| {
|
|
assert!(model.contains_any_remote_blocks());
|
|
assert!(model.contains_any_restored_remote_blocks());
|
|
});
|
|
|
|
// Verify event was emitted exactly once
|
|
let mut count = 0;
|
|
while receiver.try_recv().is_ok() {
|
|
count += 1;
|
|
}
|
|
assert_eq!(count, 1);
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_update_no_change() {
|
|
App::test((), |mut app| async move {
|
|
// Create FocusedTerminalInfo with default values (false, false)
|
|
let model_handle = app.add_model(|_| FocusedTerminalInfo::default());
|
|
|
|
// Setup event tracking
|
|
let (sender, receiver) = async_channel::unbounded();
|
|
app.update(|ctx| {
|
|
let sender = sender.clone();
|
|
ctx.subscribe_to_model(
|
|
&model_handle,
|
|
move |_, event: &FocusedTerminalInfoEvent, _| match event {
|
|
FocusedTerminalInfoEvent::TerminalInfoUpdated => {
|
|
let _ = sender.try_send(());
|
|
}
|
|
},
|
|
);
|
|
});
|
|
|
|
// First update to (true, true)
|
|
model_handle.update(&mut app, |model, ctx| {
|
|
model.update(true, true, ctx);
|
|
});
|
|
|
|
// Clear events by draining the channel
|
|
while receiver.try_recv().is_ok() {}
|
|
|
|
// Update with same values (true, true)
|
|
model_handle.update(&mut app, |model, ctx| {
|
|
model.update(true, true, ctx);
|
|
});
|
|
|
|
// Verify model state remains the same
|
|
model_handle.read(&app, |model, _| {
|
|
assert!(model.contains_any_remote_blocks());
|
|
assert!(model.contains_any_restored_remote_blocks());
|
|
});
|
|
|
|
// Verify no event was emitted
|
|
let mut count = 0;
|
|
while receiver.try_recv().is_ok() {
|
|
count += 1;
|
|
}
|
|
assert_eq!(count, 0);
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_update_only_remote_toggles() {
|
|
App::test((), |mut app| async move {
|
|
// Create FocusedTerminalInfo with default values (false, false)
|
|
let model_handle = app.add_model(|_| FocusedTerminalInfo::default());
|
|
|
|
// Setup event tracking
|
|
let (sender, receiver) = async_channel::unbounded();
|
|
app.update(|ctx| {
|
|
let sender = sender.clone();
|
|
ctx.subscribe_to_model(
|
|
&model_handle,
|
|
move |_, event: &FocusedTerminalInfoEvent, _| match event {
|
|
FocusedTerminalInfoEvent::TerminalInfoUpdated => {
|
|
let _ = sender.try_send(());
|
|
}
|
|
},
|
|
);
|
|
});
|
|
|
|
// First update to (true, true)
|
|
model_handle.update(&mut app, |model, ctx| {
|
|
model.update(true, true, ctx);
|
|
});
|
|
|
|
// Clear events by draining the channel
|
|
while receiver.try_recv().is_ok() {}
|
|
|
|
// Update with (false, true) - only remote blocks changes
|
|
model_handle.update(&mut app, |model, ctx| {
|
|
model.update(false, true, ctx);
|
|
});
|
|
|
|
// Verify model state
|
|
model_handle.read(&app, |model, _| {
|
|
assert!(!model.contains_any_remote_blocks());
|
|
assert!(model.contains_any_restored_remote_blocks());
|
|
});
|
|
|
|
// Verify event was emitted exactly once
|
|
let mut count = 0;
|
|
while receiver.try_recv().is_ok() {
|
|
count += 1;
|
|
}
|
|
assert_eq!(count, 1);
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_update_only_restored_toggles() {
|
|
App::test((), |mut app| async move {
|
|
// Create FocusedTerminalInfo with default values (false, false)
|
|
let model_handle = app.add_model(|_| FocusedTerminalInfo::default());
|
|
|
|
// Setup event tracking
|
|
let (sender, receiver) = async_channel::unbounded();
|
|
app.update(|ctx| {
|
|
let sender = sender.clone();
|
|
ctx.subscribe_to_model(
|
|
&model_handle,
|
|
move |_, event: &FocusedTerminalInfoEvent, _| match event {
|
|
FocusedTerminalInfoEvent::TerminalInfoUpdated => {
|
|
let _ = sender.try_send(());
|
|
}
|
|
},
|
|
);
|
|
});
|
|
|
|
// First update to (true, true)
|
|
model_handle.update(&mut app, |model, ctx| {
|
|
model.update(true, true, ctx);
|
|
});
|
|
|
|
// Clear events by draining the channel
|
|
while receiver.try_recv().is_ok() {}
|
|
|
|
// Update with (true, false) - only restored blocks changes
|
|
model_handle.update(&mut app, |model, ctx| {
|
|
model.update(true, false, ctx);
|
|
});
|
|
|
|
// Verify model state
|
|
model_handle.read(&app, |model, _| {
|
|
assert!(model.contains_any_remote_blocks());
|
|
assert!(!model.contains_any_restored_remote_blocks());
|
|
});
|
|
|
|
// Verify event was emitted exactly once
|
|
let mut count = 0;
|
|
while receiver.try_recv().is_ok() {
|
|
count += 1;
|
|
}
|
|
assert_eq!(count, 1);
|
|
});
|
|
}
|
|
|
|
// ToolbarCommandMap Tests
|
|
|
|
#[test]
|
|
fn test_toolbar_command_map_deserialize_from_map() {
|
|
let json = serde_json::json!({
|
|
"^claude": "Claude",
|
|
"^gemini": "Gemini",
|
|
"^codex": ""
|
|
});
|
|
let map: ToolbarCommandMap = serde_json::from_value(json).unwrap();
|
|
assert_eq!(map.0.len(), 3);
|
|
assert_eq!(map.0["^claude"], "Claude");
|
|
assert_eq!(map.0["^gemini"], "Gemini");
|
|
assert_eq!(map.0["^codex"], "");
|
|
}
|
|
|
|
#[test]
|
|
fn test_toolbar_command_map_deserialize_from_legacy_vec() {
|
|
let json = serde_json::json!(["^claude", "^gemini", "^custom"]);
|
|
let map: ToolbarCommandMap = serde_json::from_value(json).unwrap();
|
|
assert_eq!(map.0.len(), 3);
|
|
// Legacy vec format should assign empty agent values.
|
|
for (_, agent) in map.0.iter() {
|
|
assert_eq!(agent, "");
|
|
}
|
|
let keys: Vec<_> = map.0.keys().collect();
|
|
assert_eq!(keys, vec!["^claude", "^gemini", "^custom"]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_toolbar_command_map_from_file_value_map_format() {
|
|
use settings_value::SettingsValue;
|
|
|
|
let value = serde_json::json!({
|
|
"^claude": "Claude",
|
|
"^amp": "Amp"
|
|
});
|
|
let map = ToolbarCommandMap::from_file_value(&value).unwrap();
|
|
assert_eq!(map.0.len(), 2);
|
|
assert_eq!(map.0["^claude"], "Claude");
|
|
assert_eq!(map.0["^amp"], "Amp");
|
|
}
|
|
|
|
#[test]
|
|
fn test_toolbar_command_map_from_file_value_legacy_array() {
|
|
// Patterns are intentionally non-alphabetical to verify insertion order is preserved.
|
|
let value = serde_json::json!(["^zebra", "^alpha", "^middle"]);
|
|
let map = ToolbarCommandMap::from_file_value(&value).unwrap();
|
|
assert_eq!(map.0.len(), 3);
|
|
assert_eq!(map.0["^zebra"], "");
|
|
assert_eq!(map.0["^alpha"], "");
|
|
assert_eq!(map.0["^middle"], "");
|
|
let keys: Vec<_> = map.0.keys().collect();
|
|
assert_eq!(keys, vec!["^zebra", "^alpha", "^middle"]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_toolbar_command_map_from_file_value_invalid() {
|
|
let value = serde_json::json!(42);
|
|
assert!(ToolbarCommandMap::from_file_value(&value).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_toolbar_command_map_roundtrip() {
|
|
let mut inner = IndexMap::new();
|
|
inner.insert("^claude".to_string(), "Claude".to_string());
|
|
inner.insert("^custom".to_string(), String::new());
|
|
let original = ToolbarCommandMap::new(inner);
|
|
|
|
let file_value = original.to_file_value();
|
|
let restored = ToolbarCommandMap::from_file_value(&file_value).unwrap();
|
|
assert_eq!(original, restored);
|
|
}
|
|
|
|
#[test]
|
|
fn initial_litellm_provider_maps_codex_model_to_rig_without_a_committed_key() {
|
|
let providers = default_openai_providers();
|
|
|
|
assert_eq!(providers.len(), 2);
|
|
let provider = &providers[0];
|
|
assert_eq!(provider.kind, OpenAIProviderKind::LiteLLM);
|
|
assert_eq!(provider.base_url, INITIAL_LITELLM_BASE_URL);
|
|
assert_eq!(provider.api_key, None);
|
|
assert_eq!(provider.models.len(), 1);
|
|
let model = &provider.models[0];
|
|
assert_eq!(model.model_id, INITIAL_RIG_MODEL_ID);
|
|
assert!(model.use_rig);
|
|
assert_eq!(model.supports_system_messages, Some(false));
|
|
assert!(!model.supports_system_messages());
|
|
|
|
let chatgpt = &providers[1];
|
|
assert_eq!(chatgpt.kind, OpenAIProviderKind::ChatGPTSubscription);
|
|
assert_eq!(chatgpt.name, "ChatGPT Subscription");
|
|
assert!(chatgpt.base_url.is_empty());
|
|
assert!(chatgpt.api_key.is_none());
|
|
assert!(chatgpt
|
|
.models
|
|
.iter()
|
|
.any(|model| model.model_id == "gpt-5.4-pro"));
|
|
|
|
let sol = chatgpt
|
|
.models
|
|
.iter()
|
|
.find(|model| model.model_id == "gpt-5.6-sol")
|
|
.expect("GPT-5.6 Sol should be in the ChatGPT catalog");
|
|
assert_eq!(sol.reasoning_efforts.len(), 6);
|
|
assert!(sol.reasoning_efforts.iter().any(|effort| effort == "max"));
|
|
assert!(sol.reasoning_efforts.iter().any(|effort| effort == "ultra"));
|
|
|
|
let luna = chatgpt
|
|
.models
|
|
.iter()
|
|
.find(|model| model.model_id == "gpt-5.6-luna")
|
|
.expect("GPT-5.6 Luna should be in the ChatGPT catalog");
|
|
assert!(luna.reasoning_efforts.iter().any(|effort| effort == "max"));
|
|
assert!(luna
|
|
.reasoning_efforts
|
|
.iter()
|
|
.any(|effort| effort == "ultra"));
|
|
|
|
let terra = chatgpt
|
|
.models
|
|
.iter()
|
|
.find(|model| model.model_id == "gpt-5.6-terra")
|
|
.expect("GPT-5.6 Terra should be in the ChatGPT catalog");
|
|
assert!(terra.reasoning_efforts.iter().any(|effort| effort == "max"));
|
|
assert!(terra
|
|
.reasoning_efforts
|
|
.iter()
|
|
.any(|effort| effort == "ultra"));
|
|
|
|
let gpt_54 = chatgpt
|
|
.models
|
|
.iter()
|
|
.find(|model| model.model_id == "gpt-5.4")
|
|
.expect("GPT-5.4 should be in the ChatGPT catalog");
|
|
assert_eq!(
|
|
gpt_54.reasoning_efforts,
|
|
vec!["low", "medium", "high", "xhigh"]
|
|
.into_iter()
|
|
.map(str::to_string)
|
|
.collect::<Vec<_>>()
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn native_provider_settings_roundtrip_with_vertex_configuration() {
|
|
let provider: OpenAIProviderConfig = serde_json::from_value(serde_json::json!({
|
|
"kind": "vertex_ai",
|
|
"enabled": true,
|
|
"name": "Vertex production",
|
|
"base_url": "",
|
|
"project_id": "galaxy-project",
|
|
"location": "us-central1",
|
|
"models": []
|
|
}))
|
|
.expect("Vertex provider settings should deserialize");
|
|
|
|
assert_eq!(provider.kind, OpenAIProviderKind::VertexAI);
|
|
assert_eq!(provider.project_id.as_deref(), Some("galaxy-project"));
|
|
assert_eq!(provider.location.as_deref(), Some("us-central1"));
|
|
|
|
let legacy: OpenAIProviderConfig = serde_json::from_value(serde_json::json!({
|
|
"name": "Legacy provider",
|
|
"base_url": "http://localhost:4000/v1",
|
|
"models": []
|
|
}))
|
|
.expect("Legacy provider settings should remain compatible");
|
|
assert_eq!(legacy.kind, OpenAIProviderKind::LiteLLM);
|
|
assert_eq!(legacy.project_id, None);
|
|
assert_eq!(legacy.location, None);
|
|
|
|
let legacy_openai_compatible: OpenAIProviderConfig =
|
|
serde_json::from_value(serde_json::json!({
|
|
"kind": "openai_compatible",
|
|
"name": "Legacy LiteLLM provider",
|
|
"base_url": "http://localhost:4000/v1",
|
|
"models": []
|
|
}))
|
|
.expect("Legacy OpenAI-compatible provider settings should deserialize");
|
|
assert_eq!(legacy_openai_compatible.kind, OpenAIProviderKind::LiteLLM);
|
|
|
|
let native_openai: OpenAIProviderConfig = serde_json::from_value(serde_json::json!({
|
|
"kind": "openai",
|
|
"name": "OpenAI",
|
|
"base_url": "https://api.openai.com/v1",
|
|
"api_key": "sk-test",
|
|
"models": []
|
|
}))
|
|
.expect("Native OpenAI provider settings should deserialize");
|
|
assert_eq!(native_openai.kind, OpenAIProviderKind::OpenAI);
|
|
}
|
|
|
|
#[test]
|
|
fn codex_litellm_model_infers_missing_system_message_capability() {
|
|
let mut model = default_openai_providers().remove(0).models.remove(0);
|
|
model.supports_system_messages = None;
|
|
|
|
assert!(!model.supports_system_messages());
|
|
|
|
model.model_id = "gpt-4o".to_string();
|
|
assert!(model.supports_system_messages());
|
|
|
|
model.model_id = INITIAL_RIG_MODEL_ID.to_string();
|
|
model.supports_system_messages = Some(true);
|
|
assert!(model.supports_system_messages());
|
|
}
|
|
|
|
#[test]
|
|
fn test_toolbar_command_map_matched_agent() {
|
|
App::test((), |mut app| async move {
|
|
initialize_settings_for_tests(&mut app);
|
|
|
|
let mut map = IndexMap::new();
|
|
map.insert("^claude".to_string(), "Claude".to_string());
|
|
map.insert("^gemini".to_string(), "Gemini".to_string());
|
|
map.insert("^custom-tool".to_string(), String::new());
|
|
|
|
AISettings::handle(&app).update(&mut app, |settings, ctx| {
|
|
report_if_error!(settings
|
|
.cli_agent_footer_enabled_commands
|
|
.set_value(ToolbarCommandMap::new(map), ctx));
|
|
});
|
|
|
|
app.read(|ctx| {
|
|
let agent = CompiledCommandsForCodingAgentToolbar::matched_agent(ctx, "claude chat");
|
|
assert_eq!(agent, Some(CLIAgent::Claude));
|
|
|
|
let agent = CompiledCommandsForCodingAgentToolbar::matched_agent(ctx, "gemini ask");
|
|
assert_eq!(agent, Some(CLIAgent::Gemini));
|
|
|
|
let agent =
|
|
CompiledCommandsForCodingAgentToolbar::matched_agent(ctx, "custom-tool --flag");
|
|
assert_eq!(agent, Some(CLIAgent::Unknown));
|
|
|
|
let agent =
|
|
CompiledCommandsForCodingAgentToolbar::matched_agent(ctx, "unmatched-command");
|
|
assert_eq!(agent, None);
|
|
});
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn orchestration_is_enabled_when_ai_is_enabled() {
|
|
App::test((), |mut app| async move {
|
|
initialize_settings_for_tests(&mut app);
|
|
add_ai_enablement_dependencies_for_test(&mut app);
|
|
|
|
AISettings::handle(&app).read(&app, |settings, ctx| {
|
|
assert!(settings.is_orchestration_enabled(ctx));
|
|
});
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn ai_is_disabled_without_an_enabled_runtime() {
|
|
let _acp_flag = FeatureFlag::AgentClientProtocol.override_enabled(true);
|
|
App::test((), |mut app| async move {
|
|
initialize_settings_for_tests(&mut app);
|
|
add_ai_enablement_dependencies_for_test(&mut app);
|
|
|
|
AISettings::handle(&app).update(&mut app, |settings, ctx| {
|
|
settings
|
|
.bedrock_enabled
|
|
.set_value(false, ctx)
|
|
.expect("Bedrock setting should update");
|
|
settings
|
|
.openai_enabled
|
|
.set_value(false, ctx)
|
|
.expect("OpenAI setting should update");
|
|
settings
|
|
.acp_enabled
|
|
.set_value(false, ctx)
|
|
.expect("ACP setting should update");
|
|
});
|
|
|
|
AISettings::handle(&app).read(&app, |settings, ctx| {
|
|
assert!(!settings.has_enabled_ai_runtime());
|
|
assert!(!settings.is_any_ai_enabled(ctx));
|
|
});
|
|
|
|
AISettings::handle(&app).update(&mut app, |settings, ctx| {
|
|
settings
|
|
.acp_enabled
|
|
.set_value(true, ctx)
|
|
.expect("ACP setting should update");
|
|
});
|
|
|
|
AISettings::handle(&app).read(&app, |settings, ctx| {
|
|
assert!(settings.has_enabled_ai_runtime());
|
|
assert!(settings.is_any_ai_enabled(ctx));
|
|
});
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_should_display_quota_reset_banner_with_empty_history() {
|
|
App::test((), |mut app| async move {
|
|
initialize_settings_for_tests(&mut app);
|
|
|
|
AISettings::handle(&app).read(&app, |settings, _ctx| {
|
|
// With empty history, banner should not be displayed
|
|
assert!(!settings.should_display_quota_reset_banner());
|
|
});
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_should_display_quota_reset_banner_with_quota_exceeded_not_dismissed() {
|
|
App::test((), |mut app| async move {
|
|
initialize_settings_for_tests(&mut app);
|
|
|
|
// Set up a history with a previous cycle that had quota exceeded and banner not dismissed
|
|
let now = Utc::now();
|
|
let previous_end_date = now - chrono::Duration::days(15);
|
|
let current_end_date = now + chrono::Duration::days(15);
|
|
|
|
let previous_cycle = CycleInfo {
|
|
end_date: previous_end_date,
|
|
was_quota_exceeded: true,
|
|
banner_state: BannerState { dismissed: false },
|
|
};
|
|
|
|
let current_cycle = CycleInfo {
|
|
end_date: current_end_date,
|
|
was_quota_exceeded: false,
|
|
banner_state: BannerState::default(),
|
|
};
|
|
|
|
let cycle_history = vec![previous_cycle, current_cycle];
|
|
|
|
AISettings::handle(&app).update(&mut app, |settings, ctx| {
|
|
settings
|
|
.ai_request_quota_info
|
|
.set_value(AIRequestQuotaInfo { cycle_history }, ctx)
|
|
.unwrap();
|
|
});
|
|
|
|
AISettings::handle(&app).read(&app, |settings, _ctx| {
|
|
// Banner should be displayed when the previous cycle had quota exceeded and banner not dismissed
|
|
assert!(settings.should_display_quota_reset_banner());
|
|
});
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_should_display_quota_reset_banner_with_quota_exceeded_dismissed() {
|
|
App::test((), |mut app| async move {
|
|
initialize_settings_for_tests(&mut app);
|
|
|
|
// Set up a history with a previous cycle that had quota exceeded but banner was dismissed
|
|
let now = Utc::now();
|
|
let previous_end_date = now - chrono::Duration::days(15);
|
|
let current_end_date = now + chrono::Duration::days(15);
|
|
|
|
let previous_cycle = CycleInfo {
|
|
end_date: previous_end_date,
|
|
was_quota_exceeded: true,
|
|
banner_state: BannerState { dismissed: true },
|
|
};
|
|
|
|
let current_cycle = CycleInfo {
|
|
end_date: current_end_date,
|
|
was_quota_exceeded: false,
|
|
banner_state: BannerState::default(),
|
|
};
|
|
|
|
let cycle_history = vec![previous_cycle, current_cycle];
|
|
|
|
AISettings::handle(&app).update(&mut app, |settings, ctx| {
|
|
settings
|
|
.ai_request_quota_info
|
|
.set_value(AIRequestQuotaInfo { cycle_history }, ctx)
|
|
.unwrap();
|
|
});
|
|
|
|
AISettings::handle(&app).read(&app, |settings, _ctx| {
|
|
// Banner should not be displayed when the previous cycle had quota exceeded but banner was dismissed
|
|
assert!(!settings.should_display_quota_reset_banner());
|
|
});
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_should_display_quota_reset_banner_with_quota_not_exceeded() {
|
|
App::test((), |mut app| async move {
|
|
initialize_settings_for_tests(&mut app);
|
|
|
|
// Set up a history with a previous cycle that did not have quota exceeded
|
|
let now = Utc::now();
|
|
let previous_end_date = now - chrono::Duration::days(15);
|
|
let current_end_date = now + chrono::Duration::days(15);
|
|
|
|
let previous_cycle = CycleInfo {
|
|
end_date: previous_end_date,
|
|
was_quota_exceeded: false,
|
|
banner_state: BannerState::default(),
|
|
};
|
|
|
|
let current_cycle = CycleInfo {
|
|
end_date: current_end_date,
|
|
was_quota_exceeded: false,
|
|
banner_state: BannerState::default(),
|
|
};
|
|
|
|
let cycle_history = vec![previous_cycle, current_cycle];
|
|
|
|
AISettings::handle(&app).update(&mut app, |settings, ctx| {
|
|
settings
|
|
.ai_request_quota_info
|
|
.set_value(AIRequestQuotaInfo { cycle_history }, ctx)
|
|
.unwrap();
|
|
});
|
|
|
|
AISettings::handle(&app).read(&app, |settings, _ctx| {
|
|
// Banner should not be displayed when the previous cycle did not have quota exceeded
|
|
assert!(!settings.should_display_quota_reset_banner());
|
|
});
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_should_display_quota_reset_banner_with_only_one_cycle() {
|
|
App::test((), |mut app| async move {
|
|
initialize_settings_for_tests(&mut app);
|
|
|
|
// Set up a history with only one cycle
|
|
let now = Utc::now();
|
|
let current_end_date = now + chrono::Duration::days(15);
|
|
|
|
let current_cycle = CycleInfo {
|
|
end_date: current_end_date,
|
|
was_quota_exceeded: true, // Even if quota is exceeded
|
|
banner_state: BannerState::default(),
|
|
};
|
|
|
|
let cycle_history = vec![current_cycle];
|
|
|
|
AISettings::handle(&app).update(&mut app, |settings, ctx| {
|
|
settings
|
|
.ai_request_quota_info
|
|
.set_value(AIRequestQuotaInfo { cycle_history }, ctx)
|
|
.unwrap();
|
|
});
|
|
|
|
AISettings::handle(&app).read(&app, |settings, _ctx| {
|
|
// Banner should not be displayed when there's only one cycle, even if quota is exceeded
|
|
assert!(!settings.should_display_quota_reset_banner());
|
|
});
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_update_quota_info_create_new_cycle_when_none_exists() {
|
|
App::test((), |mut app| async move {
|
|
initialize_settings_for_tests(&mut app);
|
|
|
|
let now = Utc::now();
|
|
let next_refresh = now + chrono::Duration::days(30);
|
|
|
|
// Create a request limit info with quota not exceeded
|
|
let request_limit_info = create_test_request_limit_info(
|
|
100, // limit
|
|
50, // used
|
|
next_refresh,
|
|
false, // not unlimited
|
|
RequestLimitRefreshDuration::Monthly,
|
|
);
|
|
|
|
AISettings::handle(&app).update(&mut app, |settings, ctx| {
|
|
// Ensure we start with empty history
|
|
settings
|
|
.ai_request_quota_info
|
|
.set_value(
|
|
AIRequestQuotaInfo {
|
|
cycle_history: vec![],
|
|
},
|
|
ctx,
|
|
)
|
|
.unwrap();
|
|
|
|
// Update quota info
|
|
settings.update_quota_info(&request_limit_info, ctx);
|
|
});
|
|
|
|
AISettings::handle(&app).read(&app, |settings, _ctx| {
|
|
// Verify a new cycle was created
|
|
let cycle_history = &settings.ai_request_quota_info.cycle_history;
|
|
assert_eq!(cycle_history.len(), 1);
|
|
|
|
let cycle = &cycle_history[0];
|
|
assert_eq!(cycle.end_date, next_refresh);
|
|
assert!(!cycle.was_quota_exceeded);
|
|
assert!(!cycle.banner_state.dismissed);
|
|
});
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_update_quota_info_update_existing_cycle() {
|
|
App::test((), |mut app| async move {
|
|
initialize_settings_for_tests(&mut app);
|
|
|
|
let now = Utc::now();
|
|
let cycle_end_date = now + chrono::Duration::days(30);
|
|
|
|
// Set up an existing cycle
|
|
let existing_cycle = CycleInfo {
|
|
end_date: cycle_end_date,
|
|
was_quota_exceeded: false,
|
|
banner_state: BannerState::default(),
|
|
};
|
|
|
|
AISettings::handle(&app).update(&mut app, |settings, ctx| {
|
|
settings
|
|
.ai_request_quota_info
|
|
.set_value(
|
|
AIRequestQuotaInfo {
|
|
cycle_history: vec![existing_cycle],
|
|
},
|
|
ctx,
|
|
)
|
|
.unwrap();
|
|
});
|
|
|
|
// Create a request limit info with updated usage
|
|
let request_limit_info = create_test_request_limit_info(
|
|
100, // limit
|
|
75, // used (increased)
|
|
cycle_end_date,
|
|
false, // not unlimited
|
|
RequestLimitRefreshDuration::Monthly,
|
|
);
|
|
|
|
AISettings::handle(&app).update(&mut app, |settings, ctx| {
|
|
// Update quota info
|
|
settings.update_quota_info(&request_limit_info, ctx);
|
|
});
|
|
|
|
AISettings::handle(&app).read(&app, |settings, _ctx| {
|
|
// Verify the cycle was updated
|
|
let cycle_history = &settings.ai_request_quota_info.cycle_history;
|
|
assert_eq!(cycle_history.len(), 1);
|
|
|
|
let cycle = &cycle_history[0];
|
|
assert_eq!(cycle.end_date, cycle_end_date);
|
|
assert!(!cycle.was_quota_exceeded);
|
|
});
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_update_quota_info_quota_exceeded() {
|
|
App::test((), |mut app| async move {
|
|
initialize_settings_for_tests(&mut app);
|
|
|
|
let now = Utc::now();
|
|
let next_refresh = now + chrono::Duration::days(30);
|
|
|
|
// Create a request limit info with quota exceeded
|
|
let request_limit_info = create_test_request_limit_info(
|
|
100, // limit
|
|
100, // used (equal to limit, should be marked as exceeded)
|
|
next_refresh,
|
|
false, // not unlimited
|
|
RequestLimitRefreshDuration::Monthly,
|
|
);
|
|
|
|
AISettings::handle(&app).update(&mut app, |settings, ctx| {
|
|
// Update quota info
|
|
settings.update_quota_info(&request_limit_info, ctx);
|
|
});
|
|
|
|
AISettings::handle(&app).read(&app, |settings, _ctx| {
|
|
// Verify quota exceeded is set correctly
|
|
let cycle_history = &settings.ai_request_quota_info.cycle_history;
|
|
let cycle = &cycle_history[0];
|
|
assert!(cycle.was_quota_exceeded);
|
|
});
|
|
|
|
// Test with unlimited requests (should never be exceeded)
|
|
let unlimited_request_limit_info = create_test_request_limit_info(
|
|
100, // limit
|
|
200, // used (exceeds limit)
|
|
next_refresh,
|
|
true, // unlimited
|
|
RequestLimitRefreshDuration::Monthly,
|
|
);
|
|
|
|
AISettings::handle(&app).update(&mut app, |settings, ctx| {
|
|
// Update quota info
|
|
settings.update_quota_info(&unlimited_request_limit_info, ctx);
|
|
});
|
|
|
|
AISettings::handle(&app).read(&app, |settings, _ctx| {
|
|
// Verify quota exceeded is not set for unlimited plan
|
|
let cycle_history = &settings.ai_request_quota_info.cycle_history;
|
|
let cycle = &cycle_history[0];
|
|
assert!(!cycle.was_quota_exceeded);
|
|
});
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn test_mark_quota_banner_as_dismissed() {
|
|
App::test((), |mut app| async move {
|
|
initialize_settings_for_tests(&mut app);
|
|
|
|
let now = Utc::now();
|
|
|
|
// Create test cycles: two expired cycles and one future cycle
|
|
let expired_cycle_1 = CycleInfo {
|
|
end_date: now - chrono::Duration::days(30), // 30 days ago
|
|
was_quota_exceeded: true,
|
|
banner_state: BannerState { dismissed: false },
|
|
};
|
|
|
|
let expired_cycle_2 = CycleInfo {
|
|
end_date: now - chrono::Duration::days(15), // 15 days ago
|
|
was_quota_exceeded: true,
|
|
banner_state: BannerState { dismissed: false },
|
|
};
|
|
|
|
let future_cycle = CycleInfo {
|
|
end_date: now + chrono::Duration::days(15), // 15 days in future
|
|
was_quota_exceeded: false,
|
|
banner_state: BannerState { dismissed: false },
|
|
};
|
|
|
|
let cycle_history = vec![expired_cycle_1, expired_cycle_2, future_cycle];
|
|
|
|
// Set up initial state
|
|
AISettings::handle(&app).update(&mut app, |settings, ctx| {
|
|
settings
|
|
.ai_request_quota_info
|
|
.set_value(AIRequestQuotaInfo { cycle_history }, ctx)
|
|
.unwrap();
|
|
});
|
|
|
|
// Mark expired cycles as dismissed
|
|
AISettings::handle(&app).update(&mut app, |settings, ctx| {
|
|
settings.mark_quota_banner_as_dismissed(ctx);
|
|
});
|
|
|
|
// Verify the results
|
|
AISettings::handle(&app).read(&app, |settings, _ctx| {
|
|
let cycle_history = &settings.ai_request_quota_info.cycle_history;
|
|
assert_eq!(cycle_history.len(), 3);
|
|
|
|
// First cycle (oldest expired) should be dismissed
|
|
assert!(cycle_history[0].banner_state.dismissed);
|
|
// Second cycle (more recent expired) should be dismissed
|
|
assert!(cycle_history[1].banner_state.dismissed);
|
|
// Future cycle should not be dismissed
|
|
assert!(!cycle_history[2].banner_state.dismissed);
|
|
});
|
|
});
|
|
}
|