From 7e9b63112f9bbae8b62f0e3884632f88130ecf6c Mon Sep 17 00:00:00 2001
From: Frederick Muriuki Muriithi
Date: Fri, 14 Oct 2022 08:47:44 +0300
Subject: Update mock: Return correct test data.

---
 wqflask/tests/unit/base/test_trait.py | 6 ++++++
 1 file changed, 6 insertions(+)

(limited to 'wqflask')

diff --git a/wqflask/tests/unit/base/test_trait.py b/wqflask/tests/unit/base/test_trait.py
index e8c01ca9..5a27e243 100644
--- a/wqflask/tests/unit/base/test_trait.py
+++ b/wqflask/tests/unit/base/test_trait.py
@@ -48,6 +48,9 @@ class TestRetrieveTraitInfo(unittest.TestCase):
                                                        requests_mock):
         """Empty trait info"""
         conn = mock.MagicMock()
+        cursor = mock.MagicMock()
+        cursor.fetchone.return_value = {}
+        conn.cursor.return_value.__enter__.return_value = cursor
         mock_db.return_value.__enter__.return_value = conn
         requests_mock.return_value = TestNilResponse()
         with self.assertRaises(KeyError):
@@ -63,6 +66,9 @@ class TestRetrieveTraitInfo(unittest.TestCase):
         """Test that attributes are set"""
         mock_dataset = mock.MagicMock()
         conn = mock.MagicMock()
+        cursor = mock.MagicMock()
+        cursor.fetchone.return_value = [1, 2, 3, 4]
+        conn.cursor.return_value.__enter__.return_value = cursor
         mock_db.return_value.__enter__.return_value = conn
         requests_mock.return_value = TestResponse()
         type(mock_dataset).display_fields = mock.PropertyMock(
-- 
cgit v1.2.3