From 79eb434efd584abbf57ffe9e509e1427fca0df9c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?dj8yf0=CE=BCl?= <feepdake27@gmail.com>
Date: Mon, 25 Sep 2023 19:57:03 +0300
Subject: [PATCH] test: add a couple of tests for enum with discriminants set
 schema

---
 borsh/tests/test_enum_discriminants.rs | 129 +++++++++++++++++++++++++
 1 file changed, 129 insertions(+)

diff --git a/borsh/tests/test_enum_discriminants.rs b/borsh/tests/test_enum_discriminants.rs
index d8de987b6..33caa233f 100644
--- a/borsh/tests/test_enum_discriminants.rs
+++ b/borsh/tests/test_enum_discriminants.rs
@@ -8,6 +8,7 @@ extern crate alloc;
 use alloc::vec;
 
 use borsh::{from_slice, to_vec, BorshDeserialize, BorshSerialize};
+
 // sequence, no unit enums
 #[derive(BorshSerialize, BorshDeserialize, PartialEq, Eq, Clone, Copy, Debug)]
 #[borsh(use_discriminant = true)]
@@ -183,3 +184,131 @@ fn test_discriminant_serde() {
         assert_eq!(from_slice::<X>(&data).unwrap(), values[index]);
     }
 }
+
+#[cfg(feature = "unstable__schema")]
+mod schema {
+    #[cfg(not(feature = "std"))]
+    use alloc::{collections::BTreeMap, string::ToString, vec};
+
+    #[cfg(feature = "std")]
+    use std::collections::BTreeMap;
+
+    macro_rules! map(
+        () => { BTreeMap::new() };
+        { $($key:expr => $value:expr),+ } => {
+            {
+                let mut m = BTreeMap::new();
+                $(
+                    m.insert($key.to_string(), $value);
+                )+
+                m
+            }
+         };
+    );
+
+    use borsh::{
+        schema::{Definition, Fields},
+        BorshSchema,
+    };
+
+    #[allow(unused)]
+    #[derive(BorshSchema)]
+    #[borsh(use_discriminant = true)]
+    #[repr(i16)]
+    enum XY {
+        A,
+        B = -20,
+        C,
+        D(u32, u32),
+        E = 10,
+        F(u64),
+    }
+
+    #[test]
+    fn test_schema_discriminant_no_unit_type() {
+        assert_eq!("XY".to_string(), XY::declaration());
+        let mut defs = Default::default();
+        XY::add_definitions_recursively(&mut defs);
+        assert_eq!(
+            map! {
+                "XY" => Definition::Enum {
+                    tag_width: 1,
+                    variants: vec![
+                         (0, "A".to_string(), "XYA".to_string()),
+                         (-20, "B".to_string(), "XYB".to_string()),
+                         (-19, "C".to_string(), "XYC".to_string()),
+                         (-18, "D".to_string(), "XYD".to_string()),
+                         (10, "E".to_string(), "XYE".to_string()),
+                         (11, "F".to_string(), "XYF".to_string())
+                    ]
+                },
+                "XYA" => Definition::Struct{ fields: Fields::Empty },
+                "XYB" => Definition::Struct{ fields: Fields::Empty },
+                "XYC" => Definition::Struct{ fields: Fields::Empty },
+                "XYD" => Definition::Struct{ fields: Fields::UnnamedFields(
+                    vec!["u32".to_string(), "u32".to_string()]
+                )},
+                "XYE" => Definition::Struct{ fields: Fields::Empty },
+                "XYF" => Definition::Struct{ fields: Fields::UnnamedFields(
+                    vec!["u64".to_string()]
+
+                )},
+                "u32" => Definition::Primitive(4),
+                "u64" => Definition::Primitive(8)
+            },
+            defs
+        );
+    }
+
+    #[allow(unused)]
+    #[derive(BorshSchema)]
+    #[borsh(use_discriminant = false)]
+    #[repr(i16)]
+    enum XYNoDiscriminant {
+        A,
+        B = -20,
+        C,
+        D(u32, u32),
+        E = 10,
+        F(u64),
+    }
+
+    #[test]
+    fn test_schema_discriminant_no_unit_type_no_use_discriminant() {
+        assert_eq!(
+            "XYNoDiscriminant".to_string(),
+            XYNoDiscriminant::declaration()
+        );
+        let mut defs = Default::default();
+        XYNoDiscriminant::add_definitions_recursively(&mut defs);
+        assert_eq!(
+            map! {
+                "XYNoDiscriminant" => Definition::Enum {
+                    tag_width: 1,
+                    variants: vec![
+                         (0, "A".to_string(), "XYNoDiscriminantA".to_string()),
+                         (1, "B".to_string(), "XYNoDiscriminantB".to_string()),
+                         (2, "C".to_string(), "XYNoDiscriminantC".to_string()),
+                         (3, "D".to_string(), "XYNoDiscriminantD".to_string()),
+                         (4, "E".to_string(), "XYNoDiscriminantE".to_string()),
+                         (5, "F".to_string(), "XYNoDiscriminantF".to_string())
+                    ]
+                },
+                "XYNoDiscriminantA" => Definition::Struct{ fields: Fields::Empty },
+                "XYNoDiscriminantB" => Definition::Struct{ fields: Fields::Empty },
+                "XYNoDiscriminantC" => Definition::Struct{ fields: Fields::Empty },
+                "XYNoDiscriminantD" => Definition::Struct{ fields: Fields::UnnamedFields(
+                    vec!["u32".to_string(), "u32".to_string()]
+                )},
+                "XYNoDiscriminantE" => Definition::Struct{ fields: Fields::Empty },
+                "XYNoDiscriminantF" => Definition::Struct{ fields: Fields::UnnamedFields(
+                    vec!["u64".to_string()]
+
+                )},
+                "u32" => Definition::Primitive(4),
+                "u64" => Definition::Primitive(8)
+            },
+            defs
+        );
+    }
+}