Fix different call counts for both pyscopg and motor unit tests

AikidoSec · Nov 18, 2024 · 8a81865 · 8a81865
1 parent 9133c1c
commit 8a81865
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/aikido_zen/sinks/tests/motor_test.py b/aikido_zen/sinks/tests/motor_test.py
@@ -156,7 +156,7 @@ async def test_find_one_and_replace(db):
             filter=_filter, replacement={"dog_name": "test2"}
         )
 
-        called_with = mock_run_vulnerability_scan.call_args[1]
+        called_with = mock_run_vulnerability_scan.call_args_list[0][1]
         assert called_with["args"][0] == _filter
         assert called_with["op"] == "pymongo.collection.Collection.find_one_and_replace"
         assert called_with["kind"] == "nosql_injection"

diff --git a/aikido_zen/sinks/tests/psycopg_test.py b/aikido_zen/sinks/tests/psycopg_test.py
@@ -23,7 +23,7 @@ def test_cursor_execute(database_conn):
         query = "SELECT * FROM dogs"
         cursor.execute(query)
 
-        called_with = mock_run_vulnerability_scan.call_args[1]
+        called_with = mock_run_vulnerability_scan.call_args_list[0][1]
         assert called_with["args"][0] == query
         assert isinstance(called_with["args"][1], Postgres)
         assert called_with["op"] == "psycopg.Cursor.execute"