From 8f8a89d97c3dee097bc8f2284895e9404ab342b4 Mon Sep 17 00:00:00 2001
From: carlos <carducaldeira@gmail.com>
Date: Sun, 2 Jun 2024 19:33:43 -0300
Subject: [PATCH] update source/db.py for generic insertion

---
 source/config.py  |  3 ++-
 source/db.py      | 22 +++++-----------------
 tests/conftest.py |  4 +++-
 tests/test_db.py  | 25 ++++++++++++++++++-------
 4 files changed, 28 insertions(+), 26 deletions(-)

diff --git a/source/config.py b/source/config.py
index 3cfe879..90be2df 100644
--- a/source/config.py
+++ b/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"),
diff --git a/source/db.py b/source/db.py
index 6138ae8..19b19f6 100644
--- a/source/db.py
+++ b/source/db.py
@@ -5,7 +5,7 @@
 from source.config import config
 
 
-def create_database_connection():
+def create_database_connection() -> psycopg2.extensions.connection:
 
     params = config()
     print('Connecting to the postgreSQL database ...')
@@ -13,26 +13,14 @@ def create_database_connection():
 
     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
\ No newline at end of file
+    connection.close()
\ No newline at end of file
diff --git a/tests/conftest.py b/tests/conftest.py
index 825fbff..6140f4e 100644
--- a/tests/conftest.py
+++ b/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 = []
diff --git a/tests/test_db.py b/tests/test_db.py
index cd57ada..145d140 100644
--- a/tests/test_db.py
+++ b/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
\ No newline at end of file