mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
Move local ranking test helper into tests
This commit is contained in:
@@ -19,51 +19,9 @@ use aether_scheduler_core::{
|
||||
use super::candidate_affinity_cache::read_cached_scheduler_affinity_target;
|
||||
use super::candidate_resolution::EligibleLocalExecutionCandidate;
|
||||
use super::candidate_transport_ordering::{
|
||||
resolve_cached_candidate_execution_ordering, resolve_cached_transport_execution_ordering,
|
||||
CandidateExecutionOrdering,
|
||||
resolve_cached_transport_execution_ordering, CandidateExecutionOrdering,
|
||||
};
|
||||
|
||||
#[cfg(test)]
|
||||
async fn rank_local_execution_candidates(
|
||||
state: PlannerAppState<'_>,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
client_api_format: &str,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
) -> Vec<SchedulerMinimalCandidateSelectionCandidate> {
|
||||
let normalized_client_api_format = client_api_format.trim().to_ascii_lowercase();
|
||||
let ordering_config = read_scheduler_ordering_config_or_default(state).await;
|
||||
let mut candidates = candidates;
|
||||
let mut rankables = Vec::with_capacity(candidates.len());
|
||||
let mut ordering_cache = BTreeMap::new();
|
||||
|
||||
for (original_index, candidate) in candidates.iter().enumerate() {
|
||||
let ordering = resolve_cached_candidate_execution_ordering(
|
||||
state,
|
||||
&mut ordering_cache,
|
||||
candidate,
|
||||
ordering_config,
|
||||
)
|
||||
.await;
|
||||
rankables.push(rankable_candidate_from_candidate(
|
||||
candidate,
|
||||
original_index,
|
||||
ordering,
|
||||
normalized_client_api_format.as_str(),
|
||||
candidate.endpoint_api_format.as_str(),
|
||||
required_capabilities,
|
||||
false,
|
||||
));
|
||||
}
|
||||
|
||||
drop(ordering_cache);
|
||||
apply_scheduler_candidate_ranking(
|
||||
&mut candidates,
|
||||
&rankables,
|
||||
planner_ranking_context(ordering_config),
|
||||
);
|
||||
candidates
|
||||
}
|
||||
|
||||
pub(crate) async fn rank_eligible_local_execution_candidates(
|
||||
state: PlannerAppState<'_>,
|
||||
candidates: Vec<EligibleLocalExecutionCandidate>,
|
||||
@@ -224,20 +182,21 @@ async fn read_scheduler_ordering_config_or_default(
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_scheduler_core::RANKING_REASON_CACHED_AFFINITY;
|
||||
use aether_scheduler_core::{
|
||||
apply_scheduler_candidate_ranking, RANKING_REASON_CACHED_AFFINITY,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::candidate_affinity_cache::remember_scheduler_affinity_for_candidate;
|
||||
use super::{
|
||||
rank_local_execution_candidates, PlannerAppState,
|
||||
SchedulerMinimalCandidateSelectionCandidate,
|
||||
};
|
||||
use super::super::candidate_transport_ordering::resolve_cached_candidate_execution_ordering;
|
||||
use super::{PlannerAppState, SchedulerMinimalCandidateSelectionCandidate};
|
||||
use crate::ai_pipeline::planner::candidate_resolution::filter_and_rank_local_execution_candidates;
|
||||
use crate::data::auth::GatewayAuthApiKeySnapshot;
|
||||
use crate::data::GatewayDataState;
|
||||
@@ -245,6 +204,46 @@ mod tests {
|
||||
use crate::{scheduler::affinity::SCHEDULER_AFFINITY_TTL, AppState};
|
||||
use aether_data::repository::auth::StoredAuthApiKeySnapshot;
|
||||
|
||||
async fn rank_local_execution_candidates(
|
||||
state: PlannerAppState<'_>,
|
||||
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
client_api_format: &str,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
) -> Vec<SchedulerMinimalCandidateSelectionCandidate> {
|
||||
let normalized_client_api_format = client_api_format.trim().to_ascii_lowercase();
|
||||
let ordering_config = super::read_scheduler_ordering_config_or_default(state).await;
|
||||
let mut candidates = candidates;
|
||||
let mut rankables = Vec::with_capacity(candidates.len());
|
||||
let mut ordering_cache = BTreeMap::new();
|
||||
|
||||
for (original_index, candidate) in candidates.iter().enumerate() {
|
||||
let ordering = resolve_cached_candidate_execution_ordering(
|
||||
state,
|
||||
&mut ordering_cache,
|
||||
candidate,
|
||||
ordering_config,
|
||||
)
|
||||
.await;
|
||||
rankables.push(super::rankable_candidate_from_candidate(
|
||||
candidate,
|
||||
original_index,
|
||||
ordering,
|
||||
normalized_client_api_format.as_str(),
|
||||
candidate.endpoint_api_format.as_str(),
|
||||
required_capabilities,
|
||||
false,
|
||||
));
|
||||
}
|
||||
|
||||
drop(ordering_cache);
|
||||
apply_scheduler_candidate_ranking(
|
||||
&mut candidates,
|
||||
&rankables,
|
||||
super::planner_ranking_context(ordering_config),
|
||||
);
|
||||
candidates
|
||||
}
|
||||
|
||||
fn sample_candidate(
|
||||
endpoint_id: &str,
|
||||
key_id: &str,
|
||||
|
||||
@@ -454,8 +454,12 @@ fn ai_pipeline_planner_separates_local_candidate_resolution_from_ranking() {
|
||||
let candidate_ranking =
|
||||
read_workspace_file("apps/aether-gateway/src/ai_pipeline/planner/candidate_ranking.rs");
|
||||
assert!(
|
||||
candidate_ranking.contains("#[cfg(test)]\nasync fn rank_local_execution_candidates("),
|
||||
"planner/candidate_ranking.rs should keep raw local ranking as a test-only helper"
|
||||
candidate_ranking.contains("async fn rank_local_execution_candidates("),
|
||||
"planner/candidate_ranking.rs should keep raw local ranking helper inside tests"
|
||||
);
|
||||
assert!(
|
||||
!candidate_ranking.contains("#[cfg(test)]\nasync fn rank_local_execution_candidates("),
|
||||
"planner/candidate_ranking.rs should not keep the test-only ranking helper at module root"
|
||||
);
|
||||
for forbidden in [
|
||||
"struct SkippedLocalExecutionCandidate",
|
||||
|
||||
Reference in New Issue
Block a user