refactor: setup mocks for provisioner test

This commit is contained in:
2025-02-28 16:32:19 +00:00
parent e84e6afd54
commit 7ab64a94f0
6 changed files with 181 additions and 49 deletions

View File

@ -11,6 +11,7 @@ pub trait VmHostClient {
}
#[cfg(not(test))]
pub fn get_host_client(host: &VmHost, cfg: &ProvisionerConfig) -> Result<ProxmoxClient> {
Ok(match (host.kind.clone(), &cfg) {
(VmHostKind::Proxmox, ProvisionerConfig::Proxmox { qemu, ssh, .. }) => {
@ -20,3 +21,8 @@ pub fn get_host_client(host: &VmHost, cfg: &ProvisionerConfig) -> Result<Proxmox
_ => bail!("Unsupported host type"),
})
}
#[cfg(test)]
pub fn get_host_client(host: &VmHost, cfg: &ProvisionerConfig) -> Result<ProxmoxClient> {
todo!()
}

View File

@ -1,23 +1,28 @@
use crate::lightning::{AddInvoiceRequest, AddInvoiceResult, InvoiceUpdate, LightningNode};
use crate::router::{ArpEntry, Router};
use crate::settings::NetworkPolicy;
use anyhow::{anyhow, bail};
use chrono::Utc;
use chrono::{DateTime, Utc};
use fedimint_tonic_lnd::tonic::codegen::tokio_stream::Stream;
use lnvps_db::{
async_trait, IpRange, LNVpsDb, User, UserSshKey, Vm, VmCostPlan, VmHost, VmHostDisk,
VmHostKind, VmHostRegion, VmIpAssignment, VmOsImage, VmPayment, VmTemplate,
async_trait, DiskInterface, DiskType, IpRange, LNVpsDb, OsDistribution, User, UserSshKey, Vm,
VmCostPlan, VmCostPlanIntervalType, VmHost, VmHostDisk, VmHostKind, VmHostRegion,
VmIpAssignment, VmOsImage, VmPayment, VmTemplate,
};
use std::collections::HashMap;
use std::net::IpAddr;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::Mutex;
use crate::exchange::{ExchangeRateService, Ticker, TickerRate};
use crate::lightning::LightningNode;
#[derive(Debug, Clone)]
pub struct MockDb {
pub regions: Arc<Mutex<HashMap<u64, VmHostRegion>>>,
pub hosts: Arc<Mutex<HashMap<u64, VmHost>>>,
pub users: Arc<Mutex<HashMap<u64, User>>>,
pub cost_plans: Arc<Mutex<HashMap<u64, VmCostPlan>>>,
pub os_images: Arc<Mutex<HashMap<u64, VmOsImage>>>,
pub templates: Arc<Mutex<HashMap<u64, VmTemplate>>>,
pub vms: Arc<Mutex<HashMap<u64, Vm>>>,
pub ip_range: Arc<Mutex<HashMap<u64, IpRange>>>,
pub ip_assignments: Arc<Mutex<HashMap<u64, VmIpAssignment>>>,
@ -68,10 +73,57 @@ impl Default for MockDb {
api_token: "".to_string(),
},
);
let mut cost_plans = HashMap::new();
cost_plans.insert(
1,
VmCostPlan {
id: 1,
name: "mock".to_string(),
created: Utc::now(),
amount: 1,
currency: "EUR".to_string(),
interval_amount: 1,
interval_type: VmCostPlanIntervalType::Month,
},
);
let mut templates = HashMap::new();
templates.insert(
1,
VmTemplate {
id: 1,
name: "mock".to_string(),
enabled: true,
created: Utc::now(),
expires: None,
cpu: 2,
memory: 1024 * 1024 * 1024 * 2,
disk_size: 1024 * 1024 * 1024 * 64,
disk_type: DiskType::SSD,
disk_interface: DiskInterface::PCIe,
cost_plan_id: 1,
region_id: 1,
},
);
let mut os_images = HashMap::new();
os_images.insert(
1,
VmOsImage {
id: 1,
distribution: OsDistribution::Debian,
flavour: "server".to_string(),
version: "12".to_string(),
enabled: true,
release_date: Utc::now(),
url: "https://example.com/debian_12.img".to_string(),
},
);
Self {
regions: Arc::new(Mutex::new(regions)),
ip_range: Arc::new(Mutex::new(ip_ranges)),
hosts: Arc::new(Mutex::new(hosts)),
cost_plans: Arc::new(Mutex::new(cost_plans)),
templates: Arc::new(Mutex::new(templates)),
os_images: Arc::new(Mutex::new(os_images)),
users: Arc::new(Default::default()),
vms: Arc::new(Default::default()),
ip_assignments: Arc::new(Default::default()),
@ -108,7 +160,7 @@ impl LNVpsDb for MockDb {
}
async fn get_user(&self, id: u64) -> anyhow::Result<User> {
let mut users = self.users.lock().await;
let users = self.users.lock().await;
Ok(users.get(&id).ok_or(anyhow!("no user"))?.clone())
}
@ -116,9 +168,9 @@ impl LNVpsDb for MockDb {
let mut users = self.users.lock().await;
if let Some(u) = users.get_mut(&user.id) {
u.email = user.email.clone();
u.contact_email = user.contact_email.clone();
u.contact_nip17 = user.contact_nip17.clone();
u.contact_nip4 = user.contact_nip4.clone();
u.contact_email = user.contact_email;
u.contact_nip17 = user.contact_nip17;
u.contact_nip4 = user.contact_nip4;
}
Ok(())
}
@ -175,11 +227,13 @@ impl LNVpsDb for MockDb {
}
async fn get_os_image(&self, id: u64) -> anyhow::Result<VmOsImage> {
todo!()
let os_images = self.os_images.lock().await;
Ok(os_images.get(&id).ok_or(anyhow!("no image"))?.clone())
}
async fn list_os_image(&self) -> anyhow::Result<Vec<VmOsImage>> {
todo!()
let os_images = self.os_images.lock().await;
Ok(os_images.values().filter(|i| i.enabled).cloned().collect())
}
async fn get_ip_range(&self, id: u64) -> anyhow::Result<IpRange> {
@ -202,43 +256,71 @@ impl LNVpsDb for MockDb {
}
async fn get_cost_plan(&self, id: u64) -> anyhow::Result<VmCostPlan> {
todo!()
let cost_plans = self.cost_plans.lock().await;
Ok(cost_plans.get(&id).ok_or(anyhow!("no cost plan"))?.clone())
}
async fn get_vm_template(&self, id: u64) -> anyhow::Result<VmTemplate> {
todo!()
let templates = self.templates.lock().await;
Ok(templates.get(&id).ok_or(anyhow!("no template"))?.clone())
}
async fn list_vm_templates(&self) -> anyhow::Result<Vec<VmTemplate>> {
todo!()
let templates = self.templates.lock().await;
Ok(templates
.values()
.filter(|t| t.enabled && t.expires.as_ref().map(|e| *e > Utc::now()).unwrap_or(true))
.cloned()
.collect())
}
async fn list_vms(&self) -> anyhow::Result<Vec<Vm>> {
todo!()
let vms = self.vms.lock().await;
Ok(vms.values().filter(|v| !v.deleted).cloned().collect())
}
async fn list_expired_vms(&self) -> anyhow::Result<Vec<Vm>> {
todo!()
let vms = self.vms.lock().await;
Ok(vms
.values()
.filter(|v| !v.deleted && v.expires >= Utc::now())
.cloned()
.collect())
}
async fn list_user_vms(&self, id: u64) -> anyhow::Result<Vec<Vm>> {
todo!()
let vms = self.vms.lock().await;
Ok(vms
.values()
.filter(|v| !v.deleted && v.user_id == id)
.cloned()
.collect())
}
async fn get_vm(&self, vm_id: u64) -> anyhow::Result<Vm> {
todo!()
let vms = self.vms.lock().await;
Ok(vms.get(&vm_id).ok_or(anyhow!("no vm"))?.clone())
}
async fn insert_vm(&self, vm: &Vm) -> anyhow::Result<u64> {
todo!()
let mut vms = self.vms.lock().await;
let max_id = *vms.keys().max().unwrap_or(&0);
vms.insert(max_id + 1, vm.clone());
Ok(max_id + 1)
}
async fn delete_vm(&self, vm_id: u64) -> anyhow::Result<()> {
todo!()
let mut vms = self.vms.lock().await;
vms.remove(&vm_id);
Ok(())
}
async fn update_vm(&self, vm: &Vm) -> anyhow::Result<()> {
todo!()
let mut vms = self.vms.lock().await;
if let Some(v) = vms.get_mut(&vm.id) {
v.ssh_key_id = vm.ssh_key_id;
}
Ok(())
}
async fn insert_vm_ip_assignment(&self, ip_assignment: &VmIpAssignment) -> anyhow::Result<u64> {
@ -313,7 +395,8 @@ impl LNVpsDb for MockDb {
}
}
struct MockRouter {
#[derive(Debug, Clone)]
pub struct MockRouter {
pub policy: NetworkPolicy,
pub arp: Arc<Mutex<HashMap<String, ArpEntry>>>,
}
@ -358,8 +441,26 @@ impl Router for MockRouter {
#[derive(Clone, Debug, Default)]
pub struct MockNode {
invoices: Arc<Mutex<HashMap<String, MockInvoice>>>,
}
#[derive(Debug, Clone)]
struct MockInvoice {
pr: String,
expiry: DateTime<Utc>,
settle_index: u64,
}
#[async_trait]
impl LightningNode for MockNode {}
impl LightningNode for MockNode {
async fn add_invoice(&self, req: AddInvoiceRequest) -> anyhow::Result<AddInvoiceResult> {
todo!()
}
async fn subscribe_invoices(
&self,
from_payment_hash: Option<Vec<u8>>,
) -> anyhow::Result<Pin<Box<dyn Stream<Item = InvoiceUpdate> + Send>>> {
todo!()
}
}

View File

@ -7,9 +7,7 @@ use crate::host::proxmox::{
use crate::lightning::{AddInvoiceRequest, LightningNode};
use crate::provisioner::{NetworkProvisioner, Provisioner, ProvisionerMethod};
use crate::router::Router;
use crate::settings::{
NetworkAccessPolicy, NetworkPolicy, ProvisionerConfig, Settings,
};
use crate::settings::{NetworkAccessPolicy, NetworkPolicy, ProvisionerConfig, Settings};
use anyhow::{bail, Result};
use chrono::{Days, Months, Utc};
use fedimint_tonic_lnd::tonic::async_trait;
@ -53,9 +51,7 @@ impl LNVpsProvisioner {
db,
node,
rates,
router: settings
.get_router()
.expect("router config"),
router: settings.get_router().expect("router config"),
network_policy: settings.network_policy,
provisioner_config: settings.provisioner,
read_only: settings.read_only,
@ -251,10 +247,7 @@ impl Provisioner for LNVpsProvisioner {
}
// Use random network provisioner
let prov = NetworkProvisioner::new(
ProvisionerMethod::Random,
self.db.clone(),
);
let prov = NetworkProvisioner::new(ProvisionerMethod::Random, self.db.clone());
let template = self.db.get_vm_template(vm.template_id).await?;
let ip = prov.pick_ip_for_region(template.region_id).await?;
@ -461,10 +454,13 @@ mod tests {
use super::*;
use crate::exchange::DefaultRateCache;
use crate::mocks::{MockDb, MockNode};
use crate::settings::{LndConfig, ProvisionerConfig};
use crate::settings::{
ApiConfig, Credentials, LndConfig, ProvisionerConfig, QemuConfig, RouterConfig,
};
#[ignore]
#[tokio::test]
async fn test_basic_provisioner() {
async fn test_basic_provisioner() -> Result<()> {
let settings = Settings {
listen: None,
db: "".to_string(),
@ -492,13 +488,38 @@ mod tests {
},
delete_after: 0,
smtp: None,
router: None,
router: Some(RouterConfig::Mikrotik(ApiConfig {
id: "mock-router".to_string(),
url: "https://localhost".to_string(),
credentials: Credentials::UsernamePassword {
username: "admin".to_string(),
password: "password123".to_string(),
},
})),
dns: None,
nostr: None,
};
let db = Arc::new(MockDb::default());
let node = Arc::new(MockNode::default());
let rates = Arc::new(DefaultRateCache::default());
let provisioner = LNVpsProvisioner::new(settings, db, node, rates);
let provisioner = LNVpsProvisioner::new(settings, db.clone(), node.clone(), rates.clone());
let vm = db
.insert_vm(&Vm {
id: 1,
host_id: 1,
user_id: 1,
image_id: 1,
template_id: 1,
ssh_key_id: 1,
created: Utc::now(),
expires: Utc::now() + Duration::from_secs(30),
disk_id: 1,
mac_address: "00:00:00:00:00:00".to_string(),
deleted: false,
})
.await?;
provisioner.spawn_vm(1).await?;
Ok(())
}
}

View File

@ -1,4 +1,3 @@
use crate::settings::NetworkPolicy;
use anyhow::{bail, Result};
use ipnetwork::IpNetwork;
use lnvps_db::LNVpsDb;
@ -87,7 +86,7 @@ impl NetworkProvisioner {
mod tests {
use super::*;
use crate::mocks::*;
use crate::settings::NetworkAccessPolicy;
use lnvps_db::VmIpAssignment;
use std::str::FromStr;

View File

@ -40,13 +40,4 @@ pub struct ArpEntry {
#[cfg(feature = "mikrotik")]
mod mikrotik;
#[cfg(feature = "mikrotik")]
pub use mikrotik::*;
#[cfg(test)]
mod tests {
use super::*;
use crate::settings::NetworkPolicy;
#[test]
fn provision_ips_with_arp() {}
}
pub use mikrotik::*;

View File

@ -171,6 +171,7 @@ impl Settings {
Arc::new(LNVpsProvisioner::new(self.clone(), db, node, exchange))
}
#[cfg(not(test))]
pub fn get_router(&self) -> Result<Option<Arc<dyn Router>>> {
match &self.router {
Some(RouterConfig::Mikrotik(api)) => match &api.credentials {
@ -182,4 +183,17 @@ impl Settings {
_ => Ok(None),
}
}
#[cfg(test)]
pub fn get_router(&self) -> Result<Option<Arc<dyn Router>>> {
if self.router.is_some() {
let router = crate::mocks::MockRouter {
policy: self.network_policy.clone(),
arp: Arc::new(Default::default()),
};
Ok(Some(Arc::new(router)))
} else {
Ok(None)
}
}
}