Update tests for sql injection fixes

space2stats_api/tests/test_db_utils.py

@@ -1,7 +1,7 @@
 import unittest
 from unittest.mock import patch, Mock
 from app.utils.db_utils import get_summaries, get_available_fields
-
+from psycopg.sql import SQL, Identifier
 
 @patch("psycopg.connect")
 def test_get_summaries(mock_connect):
@@ -18,18 +18,21 @@ def test_get_summaries(mock_connect):
     rows, colnames = get_summaries(fields, h3_ids)
 
     mock_connect.assert_called_once()
-    mock_cursor.execute.assert_called_once_with(
-        f"""
-    SELECT hex_id, {', '.join(fields)}
-    FROM space2stats
-    WHERE hex_id IN ('hex_1')
-    """
+    sql_query = SQL(
+        """
+            SELECT {0}
+            FROM {1}
+            WHERE hex_id = ANY (%s)
+        """
+    ).format(
+        SQL(', ').join([Identifier(c) for c in ['hex_id'] + fields]),
+        Identifier("space2stats")
     )
+    mock_cursor.execute.assert_called_once_with(sql_query, [h3_ids])
 
     assert rows == [("hex_1", 100, 200)]
     assert colnames == ["hex_id", "field1", "field2"]
 
-
 @patch("psycopg.connect")
 def test_get_available_fields(mock_connect):
     mock_conn = Mock()
@@ -43,15 +46,15 @@ def test_get_available_fields(mock_connect):
 
     mock_connect.assert_called_once()
     mock_cursor.execute.assert_called_once_with(
-        f"""
+    """
     SELECT column_name
     FROM information_schema.columns
-    WHERE table_name = 'space2stats'
-    """
+    WHERE table_name = %s
+    """,
+        ["space2stats"]
     )
 
     assert columns == ["field1", "field2", "field3"]
 
-
 if __name__ == "__main__":
-    unittest.main()
+    unittest.main()