From 5085c2dc0d133c8c462a35c06da0fa2002b4466c Mon Sep 17 00:00:00 2001
From: Richard van der Hoff <richard@matrix.org>
Date: Thu, 16 Jan 2025 10:55:13 +0000
Subject: [PATCH] test: add support for generating dehydrated devices

---
 .../src/test_json/keys_query_sets.rs               | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/testing/matrix-sdk-test/src/test_json/keys_query_sets.rs b/testing/matrix-sdk-test/src/test_json/keys_query_sets.rs
index d63c84360e4..b0abfb34aa1 100644
--- a/testing/matrix-sdk-test/src/test_json/keys_query_sets.rs
+++ b/testing/matrix-sdk-test/src/test_json/keys_query_sets.rs
@@ -183,6 +183,10 @@ impl KeyQueryResponseTemplate {
             "user_id": self.user_id.clone(),
         });
 
+        if options.dehydrated {
+            device_keys["dehydrated"] = Value::Bool(true);
+        }
+
         sign_json(&mut device_keys, ed25519_secret_key, &self.user_id, device_id.as_str());
         if options.verified {
             let ssk = self
@@ -288,6 +292,7 @@ impl KeyQueryResponseTemplate {
 /// [`KeyQueryResponseTemplate`], via [`KeyQueryResponseTemplate::with_device`].
 pub struct KeyQueryResponseTemplateDeviceOptions {
     verified: bool,
+    dehydrated: bool,
 }
 
 impl KeyQueryResponseTemplateDeviceOptions {
@@ -304,6 +309,15 @@ impl KeyQueryResponseTemplateDeviceOptions {
         self.verified = verified;
         self
     }
+
+    /// Sets the option for whether the device will be marked as "dehydrated",
+    /// as per [MSC3814].
+    ///
+    /// [MSC3814]: https://github.com/matrix-org/matrix-spec-proposals/pull/3814
+    pub fn dehydrated(mut self, dehydrated: bool) -> Self {
+        self.dehydrated = dehydrated;
+        self
+    }
 }
 
 /// This set of keys/query response was generated using a local synapse.