test_tablevalidator.py

import os
import sqlite3
import tempfile
import unittest

from TAP.tablevalidator import TableValidator


def _make_db(table_names):
    """Create SQLite DBs with TAP_SCHEMA attached, mimicking real setup.

    Returns (conn, tap_schema_path) — caller is responsible for cleanup.
    """
    fd, tap_schema_path = tempfile.mkstemp(suffix='.db')
    os.close(fd)

    schema_conn = sqlite3.connect(tap_schema_path)
    schema_conn.execute('CREATE TABLE tables (table_name TEXT)')
    for name in table_names:
        schema_conn.execute('INSERT INTO tables VALUES (?)', (name,))
    schema_conn.commit()
    schema_conn.close()

    conn = sqlite3.connect(':memory:')
    conn.execute('ATTACH DATABASE ? AS TAP_SCHEMA', (tap_schema_path,))

    return conn, tap_schema_path


class TestTableValidator(unittest.TestCase):

    def setUp(self):
        self.conn, self._tap_schema_path = _make_db([
            'ps',
            'pscomppars',
            'stellarhosts',
            'TAP_SCHEMA.tables',
            'TAP_SCHEMA.columns',
            'TAP_SCHEMA.schemas',
            'cumulative',
        ])

    def tearDown(self):
        self.conn.close()
        if os.path.exists(self._tap_schema_path):
            os.unlink(self._tap_schema_path)

    def test_exact_match(self):
        v = TableValidator(self.conn)
        v.validate(['ps'])  # should not raise

    def test_case_insensitive(self):
        v = TableValidator(self.conn)
        v.validate(['PS'])
        v.validate(['Ps'])
        v.validate(['TAP_SCHEMA.Tables'])

    def test_schema_prefix_in_whitelist_bare_in_query(self):
        """TAP_SCHEMA.columns is whitelisted; query says just 'columns'."""
        v = TableValidator(self.conn)
        v.validate(['columns'])

    def test_bare_in_whitelist_unknown_schema_in_query(self):
        """'ps' is whitelisted bare; 'public.ps' has unknown schema — rejected."""
        v = TableValidator(self.conn)
        with self.assertRaises(Exception):
            v.validate(['public.ps'])

    def test_known_schema_bare_table_in_query(self):
        """'ps' is whitelisted bare; 'tap_schema.ps' uses known schema — allowed."""
        v = TableValidator(self.conn)
        v.validate(['tap_schema.ps'])

    def test_disallowed_table_raises(self):
        v = TableValidator(self.conn)
        with self.assertRaises(Exception) as ctx:
            v.validate(['pg_catalog.pg_tables'])
        self.assertIn('not available', str(ctx.exception))

    def test_multi_table_all_valid(self):
        v = TableValidator(self.conn)
        v.validate(['ps', 'pscomppars', 'stellarhosts'])

    def test_multi_table_one_invalid(self):
        v = TableValidator(self.conn)
        with self.assertRaises(Exception) as ctx:
            v.validate(['ps', 'information_schema.tables', 'stellarhosts'])
        self.assertIn('information_schema.tables', str(ctx.exception))

    def test_empty_table_list_raises(self):
        v = TableValidator(self.conn)
        with self.assertRaises(Exception):
            v.validate([])

    def test_system_catalog_blocked(self):
        v = TableValidator(self.conn)
        for bad_table in [
            'ALL_TABLES',
            'DBA_USERS',
            'V$SESSION',
            'information_schema.tables',
            'pg_catalog.pg_class',
            'EXOFOP.FILES',
        ]:
            with self.assertRaises(Exception, msg=f'{bad_table} should be blocked'):
                v.validate([bad_table])


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