based/tests/db_test.py

import unittest
from based.db import DBConnector, ColumnUpdate, ColumnCondition
from based.columns import *

conninfo = "postgresql://postgres:asarch6122@localhost"
connector = DBConnector(conninfo)


class TestDBConnector(unittest.TestCase):
    def test_connection(self):
        with connector.connection as conn:
            one: dict = conn.execute('SELECT 1 as "ONE"').fetchone()  # type: ignore
            self.assertIsNotNone(one)
            self.assertEqual(one["ONE"], 1)

    def test_columnDefinition(self):
        with connector.connection as conn:
            cdef1 = PrimarySerialColumnDefinition("id")
            cdef2 = TextColumnDefinition("name")
            cdef3 = TextColumnDefinition("name", "John Doe")
            self.assertEqual(cdef1.sql().as_string(conn), '"id" SERIAL PRIMARY KEY')
            self.assertEqual(cdef2.sql().as_string(conn), '"name" TEXT')
            self.assertEqual(
                cdef3.sql().as_string(conn), "\"name\" TEXT DEFAULT 'John Doe'"
            )

    def test_tableCreation(self):
        connector.dropTable("test_table")
        self.assertFalse(connector.tableExists("test_table"))
        connector.createTable(
            "test_table",
            [
                PrimarySerialColumnDefinition("id"),
                TextColumnDefinition("name", "John Doe"),
            ],
        )
        self.assertTrue(connector.tableExists("test_table"))
        connector.dropTable("test_table")
        self.assertFalse(connector.tableExists("test_table"))

    def test_differentColumnsWorking(self):
        connector.dropTable("test_table")
        self.assertFalse(connector.tableExists("test_table"))
        connector.createTable(
            "test_table",
            [
                PrimarySerialColumnDefinition("id"),
                TextColumnDefinition("textcol", "John Doe"),
                BooleanColumnDefinition("boolcol", True),
                TimestampColumnDefinition("tscol", date.today()),
                DoubleColumnDefinition("doublecol", 3.14),
                IntegerColumnDefinition("intcol", 100),
            ],
        )
        self.assertTrue(connector.tableExists("test_table"))
        connector.insertIntoTable("test_table", {})
        connector.insertIntoTable(
            "test_table",
            {
                "textcol": "Jane Doe",
                "boolcol": False,
                "tscol": date.today(),
                "doublecol": 3.14,
                "intcol": 100,
            },
        )
        res = connector.selectFromTable("test_table", ["*"])
        self.assertEqual(len(res), 2)
        self.assertEqual(res[0]["textcol"], "John Doe")
        self.assertEqual(res[0]["boolcol"], True)
        self.assertIsNotNone(res[0]["tscol"])
        self.assertEqual(res[0]["doublecol"], 3.14)
        self.assertEqual(res[0]["intcol"], 100)
        self.assertEqual(res[1]["textcol"], "Jane Doe")
        self.assertEqual(res[1]["boolcol"], False)
        self.assertIsNotNone(res[1]["tscol"])
        self.assertEqual(res[1]["doublecol"], 3.14)
        self.assertEqual(res[1]["intcol"], 100)
        connector.dropTable("test_table")
        self.assertFalse(connector.tableExists("test_table"))

    def test_tableInsertSelect(self):
        connector.dropTable("test_table")
        connector.createTable(
            "test_table",
            [
                PrimarySerialColumnDefinition("id"),
                TextColumnDefinition("name", "John Doe"),
            ],
        )
        connector.insertIntoTable("test_table", {"name": "John Doe"})
        connector.insertIntoTable("test_table", {"name": "Jane Doe"})
        connector.insertIntoTable("test_table", {"name": "John Smith"})
        connector.insertIntoTable("test_table", {"name": "Jane Smith"})
        rows = connector.selectFromTable("test_table", ["name"])
        self.assertEqual(len(rows), 4)
        self.assertEqual(rows[0]["name"], "John Doe")
        self.assertEqual(rows[1]["name"], "Jane Doe")
        self.assertEqual(rows[2]["name"], "John Smith")
        self.assertEqual(rows[3]["name"], "Jane Smith")
        connector.dropTable("test_table")

    def test_tableUpdate(self):
        connector.dropTable("test_table")
        connector.createTable(
            "test_table",
            [
                PrimarySerialColumnDefinition("id"),
                TextColumnDefinition("name", "John Doe"),
            ],
        )
        connector.insertIntoTable("test_table", {"name": "John Doe"})
        rows = connector.selectFromTable("test_table", ["name"])
        self.assertEqual(rows[0]["name"], "John Doe")
        connector.updateDataInTable(
            "test_table",
            [
                ColumnUpdate("name", "John"),
            ],
            [
                ColumnCondition("name", "eq", "John Doe"),
            ],
        )
        rows = connector.selectFromTable("test_table", ["name"])
        self.assertEqual(rows[0]["name"], "John")
        connector.dropTable("test_table")

    def test_tableDeleteFrom(self):
        connector.dropTable("test_table")
        connector.createTable(
            "test_table",
            [
                PrimarySerialColumnDefinition("id"),
                TextColumnDefinition("name", "John Doe"),
            ],
        )
        connector.insertIntoTable("test_table", {"name": "John Doe"})
        connector.insertIntoTable("test_table", {"name": "Jane Doe"})
        connector.insertIntoTable("test_table", {"name": "Mikhail Prokopenko"})
        rows = connector.selectFromTable("test_table", ["name"])
        self.assertEqual(len(rows), 3)
        self.assertEqual(rows[0]["name"], "John Doe")
        self.assertEqual(rows[1]["name"], "Jane Doe")
        self.assertEqual(rows[2]["name"], "Mikhail Prokopenko")
        connector.deleteFromTable(
            "test_table", [ColumnCondition("name", "eq", "John Doe")]
        )
        rows = connector.selectFromTable("test_table", ["name"])
        self.assertEqual(len(rows), 2)
        self.assertEqual(rows[0]["name"], "Jane Doe")
        self.assertEqual(rows[1]["name"], "Mikhail Prokopenko")
        connector.dropTable("test_table")


if __name__ == "__main__":
    unittest.main()