kangshifu-board/test/test_unified_target_probe.py

import importlib.util
import json
from pathlib import Path
import tempfile
import unittest

ROOT = Path(__file__).resolve().parents[1]
spec = importlib.util.spec_from_file_location('target_probe', ROOT / 'scripts/probe-unified-target.py')
probe = importlib.util.module_from_spec(spec)
spec.loader.exec_module(probe)


class Connection:
    def __init__(self, row):
        self.row = row
        self.statements = []
        self.closed = False
        self.rolled_back = False

    def execute(self, sql):
        self.statements.append(sql)
        return self

    def fetchone(self):
        return self.row

    def rollback(self):
        self.rolled_back = True

    def close(self):
        self.closed = True


class ProbeTests(unittest.TestCase):
    def setUp(self):
        self.temp = tempfile.TemporaryDirectory(prefix='ksf-target-probe-')
        self.addCleanup(self.temp.cleanup)
        self.config = Path(self.temp.name) / 'config.json'
        self.config.write_text(json.dumps({'schema': 'ksf.unified-db.v1', 'appId': 'FKieRbIoGt_main',
            'database': 'FKieRbIoGt_master', 'connection': {'service': 'ksfunified'}}))
        self.config.chmod(0o600)
        self.frozen_root = Path(self.temp.name) / 'frozen-r2'
        runtime = self.frozen_root / 'server/voc-library/unified_store.py'
        runtime.parent.mkdir(parents=True)
        runtime.write_bytes((ROOT / 'docs/unified-database-migration-20261109/release-r2/package/server/voc-library/unified_store.py').read_bytes())

    def test_mixed_report_is_exposed_without_changing_guards_or_config(self):
        before = self.config.read_bytes()
        connections, calls = [], []
        def connect(**kwargs):
            calls.append(kwargs)
            db = kwargs.get('dbname', 'FKieRbIoGt_master')
            connection = Connection((db, 'FKieRbIoGt_admin', 'FKieRbIoGt_admin', 'on', '17.9', 0))
            connections.append(connection)
            return connection
        value = probe.run(self.config, self.frozen_root, connect)
        self.assertFalse(value['targetIdentityVerified'])
        self.assertFalse(value['configMatchesPackage'])
        self.assertEqual(value['serviceIdentity']['database'], 'FKieRbIoGt_master')
        self.assertEqual(value['runtimeConnectionIdentity']['database'], 'FKieRbIoGtmaster')
        self.assertNotIn('dbname', calls[0])
        self.assertEqual(calls[1]['dbname'], 'FKieRbIoGtmaster')
        self.assertEqual(self.config.read_bytes(), before)
        for args, connection in zip(calls, connections):
            self.assertEqual(args['options'], '-c default_transaction_read_only=on')
            self.assertTrue(connection.closed and connection.rolled_back)
            self.assertEqual(len(connection.statements), 1)
            self.assertTrue(connection.statements[0].startswith('SELECT '))

    def test_corrected_runtime_matches_existing_service_identity(self):
        def connect(**kwargs):
            return Connection(('FKieRbIoGt_master', 'FKieRbIoGt_admin', 'FKieRbIoGt_admin', 'on', '17.9', 0))
        value = probe.run(self.config, ROOT, connect)
        self.assertTrue(value['targetIdentityVerified'])
        self.assertEqual(value['databaseWrites'], 0)

    def test_connection_error_does_not_expose_credentials(self):
        def fail(**kwargs):
            raise RuntimeError('postgresql://user:fixture-secret@host/database')
        value = probe.run(self.config, ROOT, fail)
        self.assertFalse(value['ok'])
        self.assertNotIn('fixture-secret', json.dumps(value))
        self.assertEqual(value['databaseWrites'], 0)

    def test_wrong_application_is_rejected_before_connection(self):
        value = json.loads(self.config.read_bytes())
        value['appId'] = 'other-app'
        self.config.write_text(json.dumps(value))
        def forbidden(**kwargs):
            self.fail('connection must not be attempted')
        with self.assertRaisesRegex(ValueError, 'CONFIG_APPLICATION_MISMATCH'):
            probe.run(self.config, ROOT, forbidden)


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