use tokio::time::Duration;
use veilid_core::{BareOpaqueRecordKey, BareRecordKey, BareSharedSecret, CryptoKind, RecordKey};
use crate::piece_leases::*;
use crate::piece_map::PieceMap;
fn create_test_peer_key() -> RecordKey {
RecordKey::new(
CryptoKind::default(),
BareRecordKey::new(
BareOpaqueRecordKey::new(&[0xaa; 32]),
Some(BareSharedSecret::new(&[0xaa; 32])),
),
)
}
#[tokio::test]
async fn test_lease_request_and_grant() {
let config = LeaseManagerConfig::default();
let mut manager = PieceLeaseManager::new(config);
let mut wanted = PieceMap::new();
wanted.set(0);
wanted.set(1);
wanted.set(2);
manager.set_wanted_pieces(&wanted).await;
let peer_key = create_test_peer_key();
let mut have_map = PieceMap::new();
have_map.set(1);
manager.add_peer(&peer_key, &have_map).await;
let lease = manager
.acquire_lease(&peer_key)
.await
.expect("Lease request should succeed");
assert_eq!(lease.piece_index, 1);
}
#[tokio::test]
async fn test_piece_completion_success() {
let config = LeaseManagerConfig::default();
let mut manager = PieceLeaseManager::new(config);
let mut wanted = PieceMap::new();
wanted.set(2); manager.set_wanted_pieces(&wanted).await;
let peer_key = create_test_peer_key();
let mut have_map = PieceMap::new();
have_map.set(2);
manager.add_peer(&peer_key, &have_map).await;
let lease = manager
.acquire_lease(&peer_key)
.await
.expect("Lease request should succeed");
let score = manager
.release_piece(lease.piece_index, CompletionResult::Success)
.await
.unwrap();
assert!(score > 1.0); }
#[tokio::test]
async fn test_piece_completion_failure() {
let config = LeaseManagerConfig::default();
let mut manager = PieceLeaseManager::new(config);
let mut wanted = PieceMap::new();
wanted.set(1); manager.set_wanted_pieces(&wanted).await;
let peer_key = create_test_peer_key();
let mut have_map = PieceMap::new();
have_map.set(1);
manager.add_peer(&peer_key, &have_map).await;
let lease = manager.acquire_lease(&peer_key).await.unwrap();
let score = manager
.release_piece(
lease.piece_index,
CompletionResult::Failure(FailureReason::VerificationFailed),
)
.await
.unwrap();
assert_eq!(manager.wanted_pieces_count().await, 1);
assert_eq!(manager.active_leases_count().await, 0);
let ranking = manager.get_peer_ranking(&peer_key).await.unwrap();
assert_eq!(ranking.successful_pieces(), 0);
assert_eq!(ranking.failed_pieces(), 1);
assert!(score < 1.0); }
#[tokio::test]
async fn test_lease_expiry() {
let mut config = LeaseManagerConfig::default();
config.lease_duration = Duration::from_millis(1);
config.cleanup_interval = Duration::from_millis(50);
let mut manager = PieceLeaseManager::new(config);
let mut wanted = PieceMap::new();
wanted.set(0);
manager.set_wanted_pieces(&wanted).await;
let peer_key = create_test_peer_key();
let mut have_map = PieceMap::new();
have_map.set(0);
manager.add_peer(&peer_key, &have_map).await;
let lease = manager.acquire_lease(&peer_key).await.unwrap();
manager
.release_piece(
lease.piece_index,
CompletionResult::Failure(FailureReason::Timeout),
)
.await
.unwrap();
assert_eq!(manager.active_leases_count().await, 0);
assert_eq!(manager.wanted_pieces_count().await, 1);
let ranking = manager.get_peer_ranking(&peer_key).await.unwrap();
assert_eq!(ranking.failed_pieces(), 1); assert!(ranking.score() < 1.0); }