From c68f6877da99db410c396e5beec077f4509e3b36 Mon Sep 17 00:00:00 2001
From: rxu17 <26471741+rxu17@users.noreply.github.com>
Date: Tue, 30 Apr 2024 23:13:27 -0700
Subject: [PATCH] fix tests

---
 tests/test_load.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/tests/test_load.py b/tests/test_load.py
index a9e8540a..f6abb7e9 100644
--- a/tests/test_load.py
+++ b/tests/test_load.py
@@ -96,10 +96,15 @@ def test_that_update_table_has_expected_calls(
 
     mock_database_ent = Mock()
     mock_database_ent.primaryKey = "PRIMARY_KEY"
-    syn.get.return_value = mock_database_ent
-    syn.tableQuery.return_value.asDataFrame.return_value = test_data
-
-    with patch.object(load, "_update_table") as patch__update_table:
+    mock_database = Mock()
+
+    with patch.object(syn, "get", return_value=mock_database_ent), patch.object(
+        syn, "tableQuery", return_value=mock_database
+    ) as patch_table_query, patch.object(
+        mock_database, "asDataFrame", return_value=test_data
+    ), patch.object(
+        load, "_update_table"
+    ) as patch__update_table:
         load.update_table(
             syn,
             databaseSynId=test_table_synid,
@@ -109,7 +114,7 @@ def test_that_update_table_has_expected_calls(
             col=cols_subset,
             toDelete=to_delete,
         )
-        syn.tableQuery.assert_called_with(
+        patch_table_query.assert_called_with(
             f"SELECT * FROM {test_table_synid} where CENTER ='test-center'"
         )