From eebdb764e125de4fba80a3fc17253a9c75b50319 Mon Sep 17 00:00:00 2001
From: Costas Drogos <costas.drogos@gmail.com>
Date: Tue, 13 Aug 2024 08:05:30 +0200
Subject: [PATCH] Add some tests for core.response.flatten_custom (#637)

Relates to #597
---
 tests/unit/test_response.py | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/tests/unit/test_response.py b/tests/unit/test_response.py
index 7f5d8a8..1b4beb5 100644
--- a/tests/unit/test_response.py
+++ b/tests/unit/test_response.py
@@ -2,7 +2,23 @@
 from unittest.mock import Mock, patch
 
 from pynetbox.core.endpoint import Endpoint
-from pynetbox.core.response import Record, RecordSet
+from pynetbox.core.response import Record, RecordSet, flatten_custom
+
+
+class FlattenCustomTest(unittest.TestCase):
+    def test_flatten_custom(self):
+        test_dicts = [
+            {"foo0": []},
+            {"foo1": [{"a": "b"}]},
+            {"foo2": [{"a": "b", "c": "d"}]},
+            {"foo3": 123},
+            {"foo4": "a"},
+            {"foo5": {"a": "b"}},
+            {"foo6": [{"a": "b", "c": "d"}]},
+        ]
+        for test_dict in test_dicts:
+            ret = flatten_custom(test_dict)
+            assert ret == test_dict
 
 
 class RecordTestCase(unittest.TestCase):