[PATCH datacenter-manager v2 20/20] tests: add example tests for remote subscription management

[Lukas Wagner <l.wagner@proxmox.com>]

This is obviously pretty incomplete, but it should demonstrate how easy
it is to write tests now, given the injected context.

Signed-off-by: Lukas Wagner <l.wagner@proxmox.com>
---
 server/src/api/nodes/tasks.rs           |   2 +-
 server/src/api/subscriptions/mod.rs     |  12 +--
 server/tests/common/environment.rs      |   9 +-
 server/tests/common/mod.rs              |  59 +++++++++--
 server/tests/common/test_application.rs |  53 +++++++++-
 server/tests/test_subscriptions.rs      | 130 ++++++++++++++++++++++++
 6 files changed, 246 insertions(+), 19 deletions(-)
 create mode 100644 server/tests/test_subscriptions.rs

diff --git a/server/src/api/nodes/tasks.rs b/server/src/api/nodes/tasks.rs
index 31ddb8f1..4c28880f 100644
--- a/server/src/api/nodes/tasks.rs
+++ b/server/src/api/nodes/tasks.rs
@@ -277,7 +277,7 @@ fn stop_task(upid: UPID, rpcenv: &mut dyn RpcEnvironment) -> Result<(), Error> {
     },
 )]
 /// Get task status.
-async fn get_task_status(upid: UPID, rpcenv: &mut dyn RpcEnvironment) -> Result<Value, Error> {
+pub async fn get_task_status(upid: UPID, rpcenv: &mut dyn RpcEnvironment) -> Result<Value, Error> {
     let auth_id: Authid = rpcenv
         .get_auth_id()
         .context("no authid available")?
diff --git a/server/src/api/subscriptions/mod.rs b/server/src/api/subscriptions/mod.rs
index 922b43f7..087dbb0c 100644
--- a/server/src/api/subscriptions/mod.rs
+++ b/server/src/api/subscriptions/mod.rs
@@ -133,7 +133,7 @@ fn key_not_found(key: &str) -> Error {
 /// additionally gated on per-remote `PRIV_RESOURCE_AUDIT` so that an operator who can audit the
 /// pool but not a specific remote does not learn which keys are pinned to it (and through that,
 /// the existence and rough size of that remote's deployment).
-fn list_keys(
+pub fn list_keys(
     rpcenv: &mut dyn RpcEnvironment,
     app: State<PdmApplication>,
 ) -> Result<Vec<SubscriptionKeyEntry>, Error> {
@@ -195,7 +195,7 @@ fn list_keys(
 ///
 /// The post-save digest is set on the response so clients can chain a follow-up mutation without
 /// a refetch round-trip.
-async fn add_keys(
+pub async fn add_keys(
     keys: Vec<String>,
     digest: Option<ConfigDigest>,
     rpcenv: &mut dyn RpcEnvironment,
@@ -277,7 +277,7 @@ async fn add_keys(
 /// Bound entries are hidden from operators who cannot audit the bound remote (mirrors the
 /// `list_keys` filter); the response is the same 404 either way so a probe cannot distinguish
 /// "key exists but you cannot see it" from "key not in pool".
-fn get_key(
+pub fn get_key(
     key: String,
     rpcenv: &mut dyn RpcEnvironment,
     app: State<PdmApplication>,
@@ -329,7 +329,7 @@ fn get_key(
 /// another admin had pinned. Refuses if the key is currently the live active key on its bound
 /// node, since dropping the pool entry would orphan that subscription on the remote: the
 /// operator must run Clear Key on the Node Subscription Status panel first.
-async fn delete_key(
+pub async fn delete_key(
     key: String,
     digest: Option<ConfigDigest>,
     rpcenv: &mut dyn RpcEnvironment,
@@ -463,7 +463,7 @@ async fn delete_key(
 /// `PRIV_SYS_MODIFY` lets the caller touch the pool config; per-remote `PRIV_RESOURCE_MODIFY`
 /// is enforced inside this handler so an operator cannot push a key to a remote they have no
 /// other authority on.
-async fn set_assignment(
+pub async fn set_assignment(
     key: String,
     remote: String,
     node: String,
@@ -1818,7 +1818,7 @@ fn compute_proposals(
 /// no longer sees. The worker itself deliberately re-reads the pool when it fires (a worker can
 /// be scheduled with delay), so a parallel admin edit between API return and worker firing is
 /// still honoured - the digest only pins the at-API-call-time plan, not the executed plan.
-async fn apply_pending(
+pub async fn apply_pending(
     digest: Option<ConfigDigest>,
     rpcenv: &mut dyn RpcEnvironment,
     app: State<PdmApplication>,
diff --git a/server/tests/common/environment.rs b/server/tests/common/environment.rs
index 5dad7b21..dbfb2918 100644
--- a/server/tests/common/environment.rs
+++ b/server/tests/common/environment.rs
@@ -1,14 +1,17 @@
 use proxmox_router::RpcEnvironment;
 
-pub struct TestRpcEnvironment;
+/// An [`RpcEnvironment`] that can be used in integration tests.
+pub struct TestRpcEnvironment {
+    pub attribs: serde_json::Value,
+}
 
 impl RpcEnvironment for TestRpcEnvironment {
     fn result_attrib_mut(&mut self) -> &mut serde_json::Value {
-        unimplemented!()
+        &mut self.attribs
     }
 
     fn result_attrib(&self) -> &serde_json::Value {
-        unimplemented!()
+        &self.attribs
     }
 
     fn env_type(&self) -> proxmox_router::RpcEnvironmentType {
diff --git a/server/tests/common/mod.rs b/server/tests/common/mod.rs
index 07a07208..93e2a9ce 100644
--- a/server/tests/common/mod.rs
+++ b/server/tests/common/mod.rs
@@ -1,9 +1,8 @@
-use std::{
-    path::{Path, PathBuf},
-    sync::{Once, OnceLock},
-};
+use std::path::{Path, PathBuf};
+use std::sync::{Once, OnceLock};
+use std::time::Duration;
 
-use anyhow::Context;
+use anyhow::{Context, Error, bail};
 use serde::de::DeserializeOwned;
 
 use proxmox_sys::fs::CreateOptions;
@@ -58,5 +57,53 @@ pub fn test_setup() {
 
 /// Create a [`TestRpcEnvironment`] for the use in a test.
 pub fn rpcenv() -> TestRpcEnvironment {
-    TestRpcEnvironment
+    TestRpcEnvironment {
+        attribs: serde_json::json!({}),
+    }
+}
+
+#[allow(unused)]
+macro_rules! assert_http_error {
+    ($result:expr, $status:expr $(, $expected:expr)? $(,)?) => {{
+        let err = $result.unwrap_err();
+        let http_err = err.downcast_ref::<proxmox_router::HttpError>().unwrap();
+
+        assert_eq!(http_err.code, $status);
+
+        $(
+            let err_text = err.to_string();
+            assert!(
+                err_text.contains($expected),
+                "'{}' not contained in '{}'",
+                $expected,
+                err_text,
+            );
+        )?
+    }};
+}
+
+#[allow(unused)]
+pub(crate) use assert_http_error;
+
+/// Wait for a task to finish.
+///
+/// At the moment, this uses a hard-coded timeout of 5 seconds.
+pub async fn wait_for_task(upid: &str) -> Result<String, Error> {
+    let upid = upid.parse::<pdm_api_types::UPID>()?;
+
+    for _ in 0..50 {
+        let response =
+            server::api::nodes::tasks::get_task_status(upid.clone(), &mut rpcenv()).await?;
+
+        if response["status"].as_str().unwrap() != "running" {
+            return Ok(response["exitstatus"]
+                .as_str()
+                .context("expected exitstatus to be a string")?
+                .into());
+        }
+
+        tokio::time::sleep(Duration::from_millis(100)).await;
+    }
+
+    bail!("worker did not finish after timeout");
 }
diff --git a/server/tests/common/test_application.rs b/server/tests/common/test_application.rs
index 892eca5b..1655a69b 100644
--- a/server/tests/common/test_application.rs
+++ b/server/tests/common/test_application.rs
@@ -1,21 +1,22 @@
 use std::collections::HashMap;
-use std::sync::Arc;
-use std::sync::Mutex;
+use std::sync::{Arc, Mutex};
 
 use anyhow::{Context, Error, bail};
 use serde::{Serialize, de::DeserializeOwned};
 
 use nix::unistd::User;
 use proxmox_client::Client;
+use proxmox_product_config::create_mocked_lock;
 use proxmox_section_config::typed::SectionConfigData;
 
 use pbs_api_types::Authid;
 use pdm_api_types::{
     ConfigDigest,
     remotes::{Remote, RemoteType},
+    subscription::{SubscriptionKeyEntry, SubscriptionKeyShadow},
 };
 
-use pdm_config::remotes::RemoteConfig;
+use pdm_config::{remotes::RemoteConfig, subscriptions::SubscriptionKeyConfig};
 
 use server::context::product_config::ProductConfig;
 use server::{
@@ -137,6 +138,12 @@ impl ContextFactory for TestApplication {
         Ok(Box::new(self.clone()))
     }
 
+    fn make_subscription_key_config(
+        &self,
+    ) -> Result<Box<dyn SubscriptionKeyConfig + Send + Sync>, Error> {
+        Ok(Box::new(TestSubscriptionKeyConfig::default()))
+    }
+
     fn make_product_config(&self) -> Result<ProductConfig, Error> {
         let user = User::from_uid(nix::unistd::getuid())
             .ok()
@@ -242,6 +249,46 @@ impl ClientFactory for TestApplication {
     }
 }
 
+#[derive(Default)]
+struct TestSubscriptionKeyConfig {
+    shadow: Mutex<SectionConfigData<SubscriptionKeyShadow>>,
+    config: Mutex<SectionConfigData<SubscriptionKeyEntry>>,
+}
+
+impl SubscriptionKeyConfig for TestSubscriptionKeyConfig {
+    fn read(&self) -> Result<(SectionConfigData<SubscriptionKeyEntry>, ConfigDigest), Error> {
+        Ok((
+            self.config.lock().unwrap().clone(),
+            ConfigDigest::from_slice([]),
+        ))
+    }
+
+    fn read_shadow(&self) -> Result<SectionConfigData<SubscriptionKeyShadow>, Error> {
+        Ok(self.shadow.lock().unwrap().clone())
+    }
+
+    fn lock(&self) -> Result<proxmox_product_config::ApiLockGuard, Error> {
+        Ok(unsafe { create_mocked_lock() })
+    }
+
+    fn write(
+        &self,
+        config: &SectionConfigData<SubscriptionKeyEntry>,
+    ) -> Result<ConfigDigest, Error> {
+        let mut guard = self.config.lock().unwrap();
+        *guard = config.clone();
+
+        Ok(ConfigDigest::from_slice([]))
+    }
+
+    fn write_shadow(&self, shadow: &SectionConfigData<SubscriptionKeyShadow>) -> Result<(), Error> {
+        let mut guard = self.shadow.lock().unwrap();
+        *guard = shadow.clone();
+
+        Ok(())
+    }
+}
+
 macro_rules! test_pve_client {
     ($ty:ident { $($overrides:tt)* }) => {
 
diff --git a/server/tests/test_subscriptions.rs b/server/tests/test_subscriptions.rs
new file mode 100644
index 00000000..122e0517
--- /dev/null
+++ b/server/tests/test_subscriptions.rs
@@ -0,0 +1,130 @@
+use http::StatusCode;
+use pdm_api_types::subscription::{ProductType, SubscriptionLevel};
+use proxmox_router::State;
+use pve_api_types::{ClusterNodeIndexResponse, SetSubscription};
+
+use crate::common::{TestApplication, TestRemoteState, rpcenv};
+
+use server::{api, context::ContextFactory};
+
+pub mod common;
+
+#[tokio::test]
+async fn test_manage_inventory() {
+    common::test_setup();
+
+    let test_app = TestApplication::new();
+    let app = test_app.make_pdm_application().unwrap();
+
+    let existing_keys = api::subscriptions::list_keys(&mut rpcenv(), State(app.clone())).unwrap();
+    assert!(existing_keys.is_empty());
+
+    let keys = vec![
+        "pve4c-aaaaaaaaaa".into(),
+        "pve4c-bbbbbbbbbb".into(),
+        "pve4c-bbbbbbbbbb".into(),
+    ];
+
+    let add_result = api::subscriptions::add_keys(keys, None, &mut rpcenv(), State(app.clone()))
+        .await
+        .unwrap();
+
+    assert_eq!(add_result.added, 2);
+    assert_eq!(add_result.deduplicated, 1);
+
+    let existing_keys = api::subscriptions::list_keys(&mut rpcenv(), State(app.clone())).unwrap();
+    assert_eq!(existing_keys.len(), 2);
+
+    let key =
+        api::subscriptions::get_key("pve4c-aaaaaaaaaa".into(), &mut rpcenv(), State(app.clone()))
+            .unwrap();
+
+    assert_eq!(key.key, "pve4c-aaaaaaaaaa");
+    assert_eq!(key.product_type, ProductType::Pve);
+    assert_eq!(key.level, SubscriptionLevel::Community);
+    assert_eq!(key.remote, None);
+
+    api::subscriptions::delete_key(
+        "pve4c-aaaaaaaaaa".into(),
+        None,
+        &mut rpcenv(),
+        State(app.clone()),
+    )
+    .await
+    .unwrap();
+
+    let result =
+        api::subscriptions::get_key("pve4c-aaaaaaaaaa".into(), &mut rpcenv(), State(app.clone()));
+
+    common::assert_http_error!(result, StatusCode::NOT_FOUND, "not found in pool");
+}
+
+#[tokio::test]
+async fn test_invalid_key_format() {
+    common::test_setup();
+
+    let test_app = TestApplication::new();
+    let app = test_app.make_pdm_application().unwrap();
+
+    let keys = vec!["aaaaa".into()];
+
+    let result = api::subscriptions::add_keys(keys, None, &mut rpcenv(), State(app.clone())).await;
+
+    common::assert_http_error!(result, StatusCode::BAD_REQUEST, "unrecognised key format");
+}
+
+common::test_pve_client!(SubscriptionClient {
+    async fn set_subscription(&self, node: &str, params: SetSubscription) -> Result<(), proxmox_client::Error> {
+        self.state().set(format!("set_subscription-{node}"), &params);
+
+        Ok(())
+    }
+});
+
+#[tokio::test]
+async fn test_apply_key() {
+    common::test_setup();
+
+    let remote_state = TestRemoteState::default();
+
+    let test_app = TestApplication::new().with_pve_remote(
+        "remote-a",
+        remote_state.clone(),
+        SubscriptionClient,
+    );
+
+    let app = test_app.make_pdm_application().unwrap();
+
+    let keys = vec![
+        "pve4c-aaaaaaaaaa".into(),
+        "pve4c-bbbbbbbbbb".into(),
+        "pve4c-bbbbbbbbbb".into(),
+    ];
+    let _ = api::subscriptions::add_keys(keys.clone(), None, &mut rpcenv(), State(app.clone()))
+        .await
+        .unwrap();
+
+    api::subscriptions::set_assignment(
+        keys[0].clone(),
+        "remote-a".into(),
+        "remote-a-node-0".into(),
+        None,
+        &mut rpcenv(),
+        State(app.clone()),
+    )
+    .await
+    .unwrap();
+
+    let upid = api::subscriptions::apply_pending(None, &mut rpcenv(), State(app.clone()))
+        .await
+        .unwrap()
+        .unwrap();
+
+    assert_eq!(common::wait_for_task(&upid).await.unwrap(), "OK");
+
+    let set_subscription: SetSubscription = remote_state
+        .get("set_subscription-remote-a-node-0")
+        .expect("subscription was set on remote-a-node-0");
+
+    assert_eq!(set_subscription.key, keys[0]);
+}
-- 
2.47.3