| from __future__ import annotations |
| |
| import sys |
| import unittest |
| from pathlib import Path |
| from unittest.mock import patch |
| |
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT / "src")) |
| |
| from iotdb_mcp_server.config import Config # noqa: E402 |
| from iotdb_mcp_server.session_manager import IoTDBSessionManager # noqa: E402 |
| from iotdb_mcp_server.services.target_selection import apply_target_registry # noqa: E402 |
| from iotdb_mcp_server.services.target_selection import _ainode_probe_cache_key # noqa: E402 |
| from iotdb_mcp_server.services.target_selection import _target_payload # noqa: E402 |
| from iotdb_mcp_server.target_registry import ( # noqa: E402 |
| IoTDBTargetRegistry, |
| mark_target_last_known_good, |
| merge_target, |
| target_from_mapping, |
| ) |
| |
| |
| def _cloud_target(password: str): |
| return target_from_mapping( |
| { |
| "target_id": "cloud-192-168-99-20", |
| "display_name": "Cloud IoTDB 192.168.99.20", |
| "kind": "cloud", |
| "host": "192.168.99.20", |
| "port": 6667, |
| "node_urls": ["192.168.99.20:6667"], |
| "user": "root", |
| "password": password, |
| "database": "test", |
| "sql_dialect": "tree", |
| "timezone": "+00:00", |
| } |
| ) |
| |
| |
| class SessionManagerAuthCacheTest(unittest.TestCase): |
| def test_auth_fingerprint_tracks_password_without_changing_public_fingerprint(self) -> None: |
| password_target = _cloud_target("TimechoDB@2021") |
| empty_password_target = _cloud_target("") |
| |
| self.assertEqual(password_target.fingerprint(), empty_password_target.fingerprint()) |
| self.assertNotEqual( |
| password_target.auth_fingerprint(), |
| empty_password_target.auth_fingerprint(), |
| ) |
| |
| def test_empty_password_override_is_preserved(self) -> None: |
| password_target = _cloud_target("TimechoDB@2021") |
| empty_password_target = merge_target(password_target, {"password": ""}) |
| |
| self.assertEqual(empty_password_target.password, "") |
| self.assertFalse(empty_password_target.password_set) |
| self.assertNotEqual( |
| password_target.auth_fingerprint(), |
| empty_password_target.auth_fingerprint(), |
| ) |
| |
| def test_last_known_good_credential_is_used_when_password_is_missing(self) -> None: |
| target = target_from_mapping( |
| { |
| "target_id": "cloud-192-168-99-20", |
| "host": "192.168.99.20", |
| "port": 6667, |
| "user": "root", |
| "last_known_good_credential": { |
| "user": "root", |
| "password": "", |
| "last_success_at": "2026-07-10T00:00:00+00:00", |
| }, |
| } |
| ) |
| |
| self.assertEqual(target.password, "") |
| self.assertEqual(target.last_known_good_credential["password"], "") |
| self.assertFalse(target.password_set) |
| |
| def test_last_known_good_non_empty_password_is_used_when_password_is_missing(self) -> None: |
| target = target_from_mapping( |
| { |
| "target_id": "cloud-192-168-99-20", |
| "host": "192.168.99.20", |
| "port": 6667, |
| "user": "root", |
| "last_known_good_credential": { |
| "user": "root", |
| "password": "known-good", |
| "last_success_at": "2026-07-10T00:00:00+00:00", |
| }, |
| } |
| ) |
| |
| self.assertEqual(target.password, "known-good") |
| self.assertTrue(target.password_set) |
| |
| def test_explicit_password_overrides_last_known_good_credential(self) -> None: |
| target = target_from_mapping( |
| { |
| "target_id": "cloud-192-168-99-20", |
| "host": "192.168.99.20", |
| "port": 6667, |
| "user": "root", |
| "password": "", |
| "last_known_good_credential": { |
| "user": "root", |
| "password": "known-good", |
| "last_success_at": "2026-07-10T00:00:00+00:00", |
| }, |
| } |
| ) |
| |
| self.assertEqual(target.password, "") |
| self.assertFalse(target.password_set) |
| |
| def test_last_known_good_credential_is_redacted_in_public_target_payload(self) -> None: |
| target = mark_target_last_known_good( |
| _cloud_target("known-good"), |
| last_success_at="2026-07-10T00:00:00+00:00", |
| ) |
| |
| public = target.as_dict(include_secret=False) |
| secret = target.as_dict(include_secret=True) |
| |
| self.assertEqual(public["password"], "***") |
| self.assertEqual(public["last_known_good_credential"]["password"], "***") |
| self.assertEqual(secret["last_known_good_credential"]["password"], "known-good") |
| |
| def test_mark_last_known_good_updates_effective_credential(self) -> None: |
| target = mark_target_last_known_good( |
| _cloud_target(""), |
| password="known-good", |
| last_success_at="2026-07-10T00:00:00+00:00", |
| ) |
| |
| self.assertEqual(target.password, "known-good") |
| self.assertEqual(target.last_known_good_credential["password"], "known-good") |
| self.assertTrue(target.password_set) |
| |
| def test_config_round_trip_preserves_last_known_good_credential(self) -> None: |
| target = mark_target_last_known_good( |
| _cloud_target(""), |
| last_success_at="2026-07-10T00:00:00+00:00", |
| ) |
| config = Config.from_target(target) |
| |
| round_trip = config.to_target() |
| |
| self.assertEqual(round_trip.password, "") |
| self.assertEqual(round_trip.last_known_good_credential["password"], "") |
| self.assertEqual( |
| round_trip.last_known_good_credential["last_success_at"], |
| "2026-07-10T00:00:00+00:00", |
| ) |
| |
| def test_tree_pool_cache_is_not_reused_across_password_change(self) -> None: |
| password_target = _cloud_target("TimechoDB@2021") |
| empty_password_target = _cloud_target("") |
| manager = IoTDBSessionManager( |
| IoTDBTargetRegistry( |
| {password_target.target_id: password_target}, |
| password_target.target_id, |
| ) |
| ) |
| created_pools = [object(), object()] |
| |
| with patch( |
| "iotdb_mcp_server.session_manager.create_tree_session_pool", |
| side_effect=created_pools, |
| ) as create_pool: |
| first_pool = manager.tree_pool(password_target.target_id) |
| manager.update_registry( |
| IoTDBTargetRegistry( |
| {empty_password_target.target_id: empty_password_target}, |
| empty_password_target.target_id, |
| ), |
| close_pools=False, |
| ) |
| second_pool = manager.tree_pool(empty_password_target.target_id) |
| |
| self.assertIs(first_pool, created_pools[0]) |
| self.assertIs(second_pool, created_pools[1]) |
| self.assertEqual(create_pool.call_count, 2) |
| |
| def test_table_pool_cache_is_not_reused_across_password_change(self) -> None: |
| password_target = _cloud_target("TimechoDB@2021") |
| empty_password_target = _cloud_target("") |
| manager = IoTDBSessionManager( |
| IoTDBTargetRegistry( |
| {password_target.target_id: password_target}, |
| password_target.target_id, |
| ) |
| ) |
| created_pools = [object(), object()] |
| |
| with patch( |
| "iotdb_mcp_server.session_manager.create_table_session_pool", |
| side_effect=created_pools, |
| ) as create_pool: |
| first_pool = manager.table_pool(password_target.target_id) |
| manager.update_registry( |
| IoTDBTargetRegistry( |
| {empty_password_target.target_id: empty_password_target}, |
| empty_password_target.target_id, |
| ), |
| close_pools=False, |
| ) |
| second_pool = manager.table_pool(empty_password_target.target_id) |
| |
| self.assertIs(first_pool, created_pools[0]) |
| self.assertIs(second_pool, created_pools[1]) |
| self.assertEqual(create_pool.call_count, 2) |
| |
| def test_ainode_probe_cache_key_tracks_password_change(self) -> None: |
| password_config = Config.from_target(_cloud_target("TimechoDB@2021")) |
| empty_password_config = Config.from_target(_cloud_target("")) |
| |
| self.assertEqual( |
| password_config.target_fingerprint, |
| empty_password_config.target_fingerprint, |
| ) |
| self.assertNotEqual( |
| _ainode_probe_cache_key(password_config), |
| _ainode_probe_cache_key(empty_password_config), |
| ) |
| |
| def test_target_payload_does_not_probe_ainode_implicitly(self) -> None: |
| config = Config.from_target(_cloud_target("known-good")) |
| |
| with patch( |
| "iotdb_mcp_server.services.target_selection._probe_ainode_availability" |
| ) as probe: |
| payload = _target_payload(config) |
| |
| probe.assert_not_called() |
| self.assertEqual(payload["ainode_available"], "unknown") |
| self.assertEqual(payload["ainode_availability_source"]["method"], "not_probed") |
| |
| def test_empty_registry_application_clears_active_connection_fields(self) -> None: |
| config = Config.from_target(_cloud_target("known-good")) |
| |
| apply_target_registry(config, IoTDBTargetRegistry({})) |
| |
| self.assertEqual(config.target_id, "") |
| self.assertEqual(config.host, "") |
| self.assertEqual(config.port, 0) |
| self.assertEqual(config.node_urls, ()) |
| self.assertEqual(config.user, "") |
| self.assertEqual(config.password, "") |
| self.assertEqual(config.last_known_good_credential, {}) |
| |
| |
| if __name__ == "__main__": |
| unittest.main() |
