added sqlite tests

Author: Arsalanjdev

tests/test_sql_db.py

@@ -1,121 +1,77 @@
 import os
-import sqlite3
+import tempfile
 
 import pytest
-from sqlalchemy import INTEGER, TEXT
-
-from rowgen.extract_from_db import DBconnect
-
-
-@pytest.fixture
-def temp_db():
-    db_path = "test.db"
-    conn = sqlite3.connect(db_path)
-    cursor = conn.cursor()
-    cursor.execute(
-        """
-    CREATE TABLE users (
-    id INTEGER PRIMARY KEY,
-    name TEXT NOT NULL,
-    username TEXT UNIQUE,
-    rank INTEGER);
-    """
+from sqlalchemy import (
+    Column,
+    Integer,
+    String,
+    ForeignKey,
+    CheckConstraint,
+    UniqueConstraint,
+    create_engine,
+    MetaData,
+    Table,
+)
+
+from rowgen.extract_from_db import (
+    extract_db_schema,
+)
+
+
+@pytest.fixture(scope="function")
+def sqlite_db_url():
+    return "sqlite:///:memory:"
+
+
+@pytest.fixture(scope="function")
+def setup_database():
+    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as tmp:
+        db_url = f"sqlite:///{tmp.name}"
+    engine = create_engine(db_url)
+    metadata = MetaData()
+
+    parent = Table(
+        "parent",
+        metadata,
+        Column("id", Integer, primary_key=True, autoincrement=True),
+        Column("name", String, nullable=False, unique=True),
+        CheckConstraint("length(name) > 1", name="name_length_check"),
     )
-    conn.commit()
-    conn.close()
-    url = f"sqlite:///{db_path}"
-    yield url
-    # Cleanup after test
-    if os.path.exists(db_path):
-        os.remove(db_path)
-
-
-@pytest.fixture
-def db_connect(temp_db):
-    dbc = DBconnect(temp_db)
-    yield dbc
-
-
-@pytest.fixture
-def empty_temp_db():
-    db_path = "empty_test.db"
-    conn = sqlite3.connect(db_path)
-    conn.close()
-    url = f"sqlite:///{db_path}"
-    yield url
-    if os.path.exists(db_path):
-        os.remove(db_path)
 
+    child = Table(
+        "child",
+        metadata,
+        Column("id", Integer, primary_key=True),
+        Column("parent_id", Integer, ForeignKey("parent.id", ondelete="CASCADE")),
+        UniqueConstraint("parent_id", name="uq_child_parent_id"),
+    )
 
-# def test_empty_db_table_columns(empty_temp_db):
-#     dbc = DBconnect(empty_temp_db)
-#     assert dbc.table_columns == {}
-
-
-def test_get_columns(db_connect):
-    # cols = db_connect.table_columns
-    expected = {
-        "users": [
-            {
-                "name": "id",
-                "type": INTEGER(),
-                "nullable": True,
-                "default": None,
-                "primary_key": 1,
-            },
-            {
-                "name": "name",
-                "type": TEXT(),
-                "nullable": False,
-                "default": None,
-                "primary_key": 0,
-            },
-            {
-                "name": "username",
-                "type": TEXT(),
-                "nullable": True,
-                "default": None,
-                "primary_key": 0,
-            },
-            {
-                "name": "rank",
-                "type": INTEGER(),
-                "nullable": True,
-                "default": None,
-                "primary_key": 0,
-            },
-        ]
-    }
-
-    result = db_connect.table_columns
-
-    # We compare only the relevant parts because type() instances won't compare cleanly
-    def clean(col):
-        return {
-            "name": col["name"],
-            "type": type(col["type"]),  # type comparison by class
-            "nullable": col["nullable"],
-            "default": col["default"],
-            "primary_key": col["primary_key"],
-        }
-
-    cleaned_result = {
-        table: [clean(col) for col in cols] for table, cols in result.items()
-    }
+    metadata.create_all(engine)
+    yield db_url  # yield the URL instead of the engine
+    engine.dispose()
+    os.remove(tmp.name)
 
-    cleaned_expected = {
-        table: [clean(col) for col in cols] for table, cols in expected.items()
-    }
 
-    assert cleaned_result == cleaned_expected
+def test_extract_schema_basic(setup_database):
+    db_url = setup_database
+    schema = extract_db_schema(db_url)
+    assert "tables" in schema
+    assert "parent" in schema["tables"]
+    assert "child" in schema["tables"]
 
+    parent = schema["tables"]["parent"]
+    assert any(c["primary_key"] for c in parent["columns"])
+    assert any(c["unique"] or c["name"] == "name" for c in parent["columns"])
+    assert len(parent["check_constraints"]) > 0
+    assert parent["primary_key"] == ["id"]
 
-# def test_invalid_db_url():
-#     bad_url = "sqlite:///non_existent_folder/non_existent_file.db"
-#     with pytest.raises(OperationalError):
-#         DBconnect(bad_url)
+    child = schema["tables"]["child"]
+    assert any(fk["target_table"] == "parent" for fk in child["foreign_keys"])
+    assert child["primary_key"] == ["id"]
+    assert any(uc["name"] == "uq_child_parent_id" for uc in child["unique_constraints"])
 
 
-def test_tables_list(db_connect):
-    tables = db_connect.tables
-    assert "users" in tables
+def test_invalid_url_raises():
+    with pytest.raises(Exception, match="Failed to extract schema"):
+        extract_db_schema("invalid_url")