refactor: setup mocks for provisioner test

src/host/mod.rs

@@ -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!()
+}

src/mocks.rs

@@ -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,
+    async_trait, DiskInterface, DiskType, IpRange, LNVpsDb, OsDistribution, User, UserSshKey, Vm,
-    VmHostKind, VmHostRegion, VmIpAssignment, VmOsImage, VmPayment, VmTemplate,
+    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_email = user.contact_email;
-            u.contact_nip17 = user.contact_nip17.clone();
+            u.contact_nip17 = user.contact_nip17;
-            u.contact_nip4 = user.contact_nip4.clone();
+            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!()
+    }
+}

src/provisioner/lnvps.rs

@@ -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::{
+use crate::settings::{NetworkAccessPolicy, NetworkPolicy, ProvisionerConfig, 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
+            router: settings.get_router().expect("router config"),
-                .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(
+        let prov = NetworkProvisioner::new(ProvisionerMethod::Random, self.db.clone());
-            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(())
     }
 }

src/provisioner/network.rs

@@ -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;
 

src/router/mod.rs

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

src/settings.rs

@@ -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)
+        }
+    }
 }