update source/db.py for generic insertion

source/config.py

@@ -1,7 +1,8 @@
 import os
+from typing import Dict
 
 
-def config():
+def config() -> Dict[str,str]:
 
     dict_param ={'dbname': os.getenv("DB_NAME"),
                 'user': os.getenv("DB_USER"),

source/db.py

@@ -5,34 +5,22 @@
 from source.config import config
 
 
-def create_database_connection():
+def create_database_connection() -> psycopg2.extensions.connection:
 
     params = config()
     print('Connecting to the postgreSQL database ...')
     connection = psycopg2.connect(**params)
 
     return connection
 
-def insert_list_database(data_list: List[Tuple]):
+def insert_list_database(sql_query: str, data_list: List[Tuple]) -> None:
 
     connection = create_database_connection()
     cursor = connection.cursor()
-    inserted_ids = []
-
-    for data in data_list:
-        cursor.execute("""
-            INSERT INTO raw_db.news (
-                author, title, description, url, image_url, publication_date,
-                content, tags, source, query0, query1, query2
-            ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) 
-        RETURNING id;
-        """, data)
-        inserted_id = cursor.fetchone()[0]
-        inserted_ids.append(inserted_id)
+
+    cursor.executemany(sql_query, data_list)
 
     connection.commit()
 
     cursor.close()
-    connection.close()
-
-    return inserted_ids
+    connection.close()

tests/conftest.py

@@ -1,5 +1,7 @@
 import pytest
 from faker import Faker
+from typing import List, Tuple
+
 
 from source.db import create_database_connection
 
@@ -14,7 +16,7 @@ def db_connection():
 
 #scope = function inicializa create_fake_news para para cada função testado
 @pytest.fixture(scope='function')
-def create_fake_news():
+def create_fake_news() -> List[Tuple]:
 
     lists = []
     news = []

tests/test_db.py

@@ -7,14 +7,25 @@
 def test_insert_database(db_connection, create_fake_news):
 
     fake_list = create_fake_news
-    id = insert_list_database(fake_list)[0]
+
+    sql_command = """
+        INSERT INTO raw_db.news (
+            author, title, description, url, image_url, publication_date,
+            content, tags, source, query0, query1, query2
+        ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s);
+    """
+
+    insert_list_database(sql_command, fake_list)
 
     cursor = db_connection.cursor()
-    cursor.execute("""SELECT author, title, description, 
-                   url, image_url, publication_date, 
-                   content, tags, source, query0, query1, query2 
-                   FROM raw_db.news WHERE id = %s;""", (id,))
-    result = cursor.fetchall()
+    sql_select = "SELECT id FROM raw_db.news WHERE title IN %s;"
+
+    titles = tuple([fake_list[0][1]])
+
+    # Executando o comando 'SELECT'
+    cursor.execute(sql_select, (titles,))
+    id = cursor.fetchone()[0]
+
     cursor.close()
 
-    assert result == fake_list
+    assert id is not None