aboutsummaryrefslogtreecommitdiff
path: root/wqflask/tests/unit/base/test_trait.py
diff options
context:
space:
mode:
Diffstat (limited to 'wqflask/tests/unit/base/test_trait.py')
-rw-r--r--wqflask/tests/unit/base/test_trait.py6
1 files changed, 6 insertions, 0 deletions
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(