#!/usr/bin/env python3 """ Unit tests for keyspace_shard_stats.py. Run: python3 scripts/keyspace_shard_stats/test_keyspace_shard_stats.py """ import os import sys import unittest from unittest.mock import MagicMock, patch sys.path.insert(0, os.path.join(os.path.dirname(__file__))) import keyspace_shard_stats as kss def make_region(region_id, start_hex, end_hex, leader_store_id=1): return { "id": region_id, "start_key": start_hex, "end_key": end_hex, "leader": {"id": region_id * 10, "store_id": leader_store_id}, "peers": [{"id": region_id * 10, "store_id": leader_store_id}], } def make_store_response(store_id, status_address): return { "store": { "id": store_id, "address": f"127.0.0.1:{10000 + store_id}", "status_address": status_address, }, "status": {}, } def make_shard_stats(region_id, keyspace_id, **overrides): stats = { "id": region_id, "ver": 7, "keyspace": keyspace_id, "active": True, "compacting": False, "flushed": True, "mem-table-count": 1, "mem-table-size": 10, "mem-table-unpersisted-props-size": 2, "l0-table-count": 3, "blob-table-count": 4, "in-use-blob-size": 5, "total-blob-size": 6, "l0-table-size": 7, "l0-cf-table-size": [1, 2, 3], "cfs": [ { "levels": [ { "level": 1, "num-tables": 2, "data-size": 11, "index-size": 12, "filter-size": 13, "max-ts": 14, "entries": 15, "old-entries": 16, "tombs": 17, "kv-size": 18, "in-use-blob-size": 19, } ] } ], "index-size": 20, "filter-size": 21, "max-ts": 22, "entries": 23, "old-entries": 24, "old-data-size": 25, "tombs": 26, "kv-size": 27, "keyspace-prefix-tables": 1, "base-version": 28, "meta-sequence": 29, "write-sequence": 30, "total-size": 31, "partial-l0s": 1, "shared-blob-tables": 2, "partial-tbls": 3, "compaction-score": 1.5, "has-over-bound-data": True, "has-del-prefixes": True, "ready-to-destroy-range": True, "trim-over-bound": False, "manual-major-compaction": True, "storage-class-spec": "IA", "ia": { "num-shards": 1, "num-tables": 2, "data-size": 3, "kv-size": 4, }, "is-sync": True, "checked-schema-version": 31, "txn-file-locks": 32, "schema-version": 33, "schema-restore-version": 34, "columnar-tables": 35, "unconverted-l0-count": 36, "columnar-levels": [{"num-files": 2, "data-size": 40, "data-kv-size": 41}], "fts-levels": [{"num-files": 3, "data-size": 42, "pk-total": 43}], "vector-indexes": {"num-files": 4, "data-size": 44}, "encryption": True, } stats.update(overrides) return stats class TestKeyspaceEncoding(unittest.TestCase): def test_encode_txn_keyspace_range(self): start, end = kss.encode_txn_keyspace_range(10) self.assertEqual(start.hex(), "7800000a") self.assertEqual(end.hex(), "7800000b") def test_encode_txn_keyspace_pd_range(self): start, end = kss.encode_txn_keyspace_pd_range(10) self.assertEqual(start.hex(), "7800000a00000000fb") self.assertEqual(end.hex(), "7800000b00000000fb") def test_encode_large_keyspace(self): start, end = kss.encode_txn_keyspace_range(100000) self.assertEqual(start.hex(), "780186a0") self.assertEqual(end.hex(), "780186a1") def test_rejects_unencodable_keyspace(self): with self.assertRaises(ValueError): kss.encode_txn_keyspace_range(0x00FFFFFF) def test_encode_table_prefix_key(self): self.assertEqual( kss.encode_table_prefix_key(123).hex(), "74800000000000007b", ) def test_encode_table_range(self): start, end = kss.encode_table_range(2, 123) self.assertEqual(start.hex(), "7800000274800000000000007b") self.assertEqual(end.hex(), "7800000274800000000000007c") def test_encode_index_prefix_key(self): self.assertEqual( kss.encode_index_prefix_key(123, 5).hex(), "74800000000000007b5f698000000000000005", ) def test_encode_index_range(self): start, end = kss.encode_index_range(2, 123, 5) self.assertEqual(start.hex(), "7800000274800000000000007b5f698000000000000005") self.assertEqual(end.hex(), "7800000274800000000000007b5f698000000000000006") def test_rejects_i64_max_table_range(self): with self.assertRaises(ValueError): kss.encode_table_range(2, kss.I64_MAX) def test_build_table_scope_keys(self): scope, start, end, pd_start, pd_end = kss.build_scope_keys(2, 123) self.assertEqual(scope, "table") self.assertEqual(start.hex(), "7800000274800000000000007b") self.assertEqual(end.hex(), "7800000274800000000000007c") self.assertEqual(pd_start.hex(), "7800000274800000ff000000007b000000fc") self.assertEqual(pd_end.hex(), "7800000274800000ff000000007c000000fc") def test_build_index_scope_keys(self): scope, start, end, pd_start, pd_end = kss.build_scope_keys(2, 123, 5) self.assertEqual(scope, "index") self.assertEqual(start.hex(), "7800000274800000000000007b5f698000000000000005") self.assertEqual(end.hex(), "7800000274800000000000007b5f698000000000000006") self.assertEqual( pd_start.hex(), "7800000274800000ff000000007b5f6980ff0000000000000500fe", ) self.assertEqual( pd_end.hex(), "7800000274800000ff000000007b5f6980ff0000000000000600fe", ) class TestHttpHelpers(unittest.TestCase): def test_normalize_status_address(self): self.assertEqual( kss.normalize_status_address("127.0.0.1:20180", False), "http://127.0.0.1:20180", ) self.assertEqual( kss.normalize_status_address("127.0.0.1:20180", True), "https://127.0.0.1:20180", ) self.assertEqual( kss.normalize_status_address("http://127.0.0.1:20180/", True), "http://127.0.0.1:20180", ) @patch("keyspace_shard_stats.pd_request") def test_scan_regions_batch_uses_hex_format(self, mock_pd): mock_pd.return_value = {"regions": [make_region(1, "7800000a", "7800000b")]} regions = kss.scan_regions_batch( ["http://pd:2379"], bytes.fromhex("7800000a"), bytes.fromhex("7800000b"), 16, None, 10, ) self.assertEqual(len(regions), 1) params = mock_pd.call_args[1]["params"] self.assertEqual(params["key"], "7800000a") self.assertEqual(params["end_key"], "7800000b") self.assertEqual(params["limit"], "16") self.assertEqual(params["format"], "hex") @patch("keyspace_shard_stats.pd_request") def test_store_status_address_cache(self, mock_pd): mock_pd.return_value = make_store_response(1, "127.0.0.1:20180") cache = {} first = kss.get_store_status_address(["http://pd:2379"], 1, None, 10, False, cache) second = kss.get_store_status_address(["http://pd:2379"], 1, None, 10, False, cache) self.assertEqual(first, "http://127.0.0.1:20180") self.assertEqual(second, first) self.assertEqual(mock_pd.call_count, 1) class TestAggregation(unittest.TestCase): def test_aggregate_two_shards(self): aggregate = kss.new_aggregate_stats() kss.aggregate_shard_stats(aggregate, make_shard_stats(101, 10)) kss.aggregate_shard_stats( aggregate, make_shard_stats( 102, 10, active=False, **{ "mem-table-size": 90, "max-ts": 100, "compaction-score": 0.5, "ready-to-destroy-range": False, "storage-class-spec": "", }, ), ) self.assertEqual(aggregate["num-shards"], 2) self.assertEqual(aggregate["num-active-shards"], 1) self.assertEqual(aggregate["num-pending-compaction-shards"], 1) self.assertEqual(aggregate["mem-table-size"], 100) self.assertEqual(aggregate["max-max-ts"], 100) self.assertEqual(aggregate["l0-cf-table-size"], [2, 4, 6]) self.assertEqual(aggregate["cfs"][0]["levels"][0]["num-tables"], 4) self.assertEqual(aggregate["columnar-levels"][0]["data-size"], 80) self.assertEqual(aggregate["fts-levels"][0]["pk-total"], 86) self.assertEqual(aggregate["vector-indexes"]["data-size"], 88) self.assertEqual(aggregate["ia"]["num-tables"], 4) self.assertEqual(aggregate["storage-class-spec-counts"]["IA"], 1) self.assertEqual(len(aggregate["ready-destroy-range-shards"]), 1) class TestCollectKeyspaceStats(unittest.TestCase): def _args(self, **kwargs): defaults = dict( keyspace_id=10, table_id=None, index_id=None, pd="http://pd:2379", region_limit=128, scan_interval=0, timeout=10, ca_cert="", client_cert="", client_key="", include_regions=False, allow_partial=False, ) defaults.update(kwargs) return MagicMock(**defaults) @patch("keyspace_shard_stats.get_region_shard_stats") @patch("keyspace_shard_stats.pd_request") def test_collects_from_region_leaders(self, mock_pd, mock_stats): region1 = make_region(101, "7800000a00000000fb", "7800000a80000000fb", 1) region2 = make_region(102, "7800000a80000000fb", "7800000b00000000fb", 1) mock_pd.side_effect = [ {"regions": [region1, region2]}, make_store_response(1, "127.0.0.1:20180"), ] mock_stats.side_effect = [ make_shard_stats(101, 10), make_shard_stats(102, 10, **{"mem-table-size": 90}), ] result = kss.collect_keyspace_stats( self._args(include_regions=True), ["http://pd:2379"], None ) self.assertTrue(result["complete"]) self.assertEqual(result["regions"]["seen"], 2) self.assertEqual(result["regions"]["summed"], 2) self.assertEqual(result["stores"]["1"]["regions"], 2) self.assertEqual(result["stats"]["num-shards"], 2) self.assertEqual(result["stats"]["mem-table-size"], 100) self.assertEqual(result["pd-start-key"], "7800000a00000000fb") self.assertEqual(result["pd-end-key"], "7800000b00000000fb") self.assertTrue(result["range-exact-by-region-boundaries"]) self.assertEqual(len(result["region-details"]), 2) self.assertEqual(mock_pd.call_count, 2) @patch("keyspace_shard_stats.get_region_shard_stats") @patch("keyspace_shard_stats.pd_request") def test_collects_table_scope(self, mock_pd, mock_stats): pd_start = "7800000274800000ff000000007b000000fc" pd_end = "7800000274800000ff000000007c000000fc" region = make_region(101, pd_start, pd_end, 1) mock_pd.side_effect = [ {"regions": [region]}, make_store_response(1, "127.0.0.1:20180"), ] mock_stats.return_value = make_shard_stats(101, 2) result = kss.collect_keyspace_stats( self._args(keyspace_id=2, table_id=123, include_regions=True), ["http://pd:2379"], None, ) self.assertTrue(result["complete"]) self.assertEqual(result["scope"], "table") self.assertEqual(result["table-id"], 123) self.assertEqual(result["regions"]["seen"], 1) self.assertEqual(result["regions"]["range-boundary-overlaps"], 0) self.assertTrue(result["range-exact-by-region-boundaries"]) self.assertEqual(result["pd-start-key"], pd_start) self.assertEqual(result["pd-end-key"], pd_end) self.assertFalse(result["region-details"][0]["range-boundary-overlap"]) @patch("keyspace_shard_stats.get_region_shard_stats") @patch("keyspace_shard_stats.pd_request") def test_collects_index_scope(self, mock_pd, mock_stats): pd_start = "7800000274800000ff000000007b5f6980ff0000000000000500fe" pd_end = "7800000274800000ff000000007b5f6980ff0000000000000600fe" region = make_region(101, pd_start, pd_end, 1) mock_pd.side_effect = [ {"regions": [region]}, make_store_response(1, "127.0.0.1:20180"), ] mock_stats.return_value = make_shard_stats(101, 2) result = kss.collect_keyspace_stats( self._args(keyspace_id=2, table_id=123, index_id=5, include_regions=True), ["http://pd:2379"], None, ) self.assertTrue(result["complete"]) self.assertEqual(result["scope"], "index") self.assertEqual(result["table-id"], 123) self.assertEqual(result["index-id"], 5) self.assertEqual(result["regions"]["seen"], 1) self.assertEqual(result["regions"]["range-boundary-overlaps"], 0) self.assertTrue(result["range-exact-by-region-boundaries"]) self.assertEqual(result["pd-start-key"], pd_start) self.assertEqual(result["pd-end-key"], pd_end) self.assertFalse(result["region-details"][0]["range-boundary-overlap"]) @patch("keyspace_shard_stats.get_region_shard_stats") @patch("keyspace_shard_stats.pd_request") def test_collects_table_scope_marks_partial_region(self, mock_pd, mock_stats): region = make_region( 101, "7800000200000000fb", "7800000300000000fb", 1, ) mock_pd.side_effect = [ {"regions": [region]}, make_store_response(1, "127.0.0.1:20180"), ] mock_stats.return_value = make_shard_stats(101, 2) result = kss.collect_keyspace_stats( self._args(keyspace_id=2, table_id=123), ["http://pd:2379"], None ) self.assertTrue(result["complete"]) self.assertEqual(result["regions"]["range-boundary-overlaps"], 1) self.assertFalse(result["range-exact-by-region-boundaries"]) @patch("keyspace_shard_stats.pd_request") def test_missing_leader_marks_partial(self, mock_pd): mock_pd.return_value = { "regions": [{"id": 101, "start_key": "7800000a", "end_key": "7800000b"}] } result = kss.collect_keyspace_stats(self._args(), ["http://pd:2379"], None) self.assertFalse(result["complete"]) self.assertEqual(result["regions"]["seen"], 1) self.assertEqual(result["regions"]["without-leader"], 1) self.assertEqual(result["stats"]["num-shards"], 0) @patch("keyspace_shard_stats.get_region_shard_stats") @patch("keyspace_shard_stats.pd_request") def test_mismatched_stats_are_skipped(self, mock_pd, mock_stats): region = make_region(101, "7800000a", "7800000b", 1) mock_pd.side_effect = [ {"regions": [region]}, make_store_response(1, "127.0.0.1:20180"), ] mock_stats.return_value = make_shard_stats(999, 10) result = kss.collect_keyspace_stats(self._args(), ["http://pd:2379"], None) self.assertFalse(result["complete"]) self.assertEqual(result["regions"]["mismatched-stats"], 1) self.assertEqual(result["regions"]["summed"], 0) @patch("keyspace_shard_stats.pd_request") def test_non_advancing_region_is_fatal(self, mock_pd): region = make_region(101, "7800000a", "7800000a", 1) mock_pd.return_value = {"regions": [region]} result = kss.collect_keyspace_stats(self._args(region_limit=1), ["http://pd:2379"], None) self.assertFalse(result["complete"]) self.assertIn("did not advance", result["fatal-error"]) class TestParseArgs(unittest.TestCase): def test_defaults(self): args = kss.parse_args(["10"]) self.assertEqual(args.keyspace_id, 10) self.assertIsNone(args.table_id) self.assertIsNone(args.index_id) self.assertEqual(args.pd, "http://127.0.0.1:2379") self.assertEqual(args.region_limit, kss.DEFAULT_REGION_LIMIT) self.assertFalse(args.include_regions) self.assertFalse(args.allow_partial) def test_optional_positional_table_id(self): args = kss.parse_args(["10", "123"]) self.assertEqual(args.keyspace_id, 10) self.assertEqual(args.table_id, 123) def test_table_id_option(self): args = kss.parse_args(["10", "--table-id", "123"]) self.assertEqual(args.keyspace_id, 10) self.assertEqual(args.table_id, 123) def test_index_id_option(self): args = kss.parse_args(["10", "--table-id", "123", "--index-id", "5"]) self.assertEqual(args.keyspace_id, 10) self.assertEqual(args.table_id, 123) self.assertEqual(args.index_id, 5) def test_positional_table_id_with_index_id(self): args = kss.parse_args(["10", "123", "--index-id", "5"]) self.assertEqual(args.keyspace_id, 10) self.assertEqual(args.table_id, 123) self.assertEqual(args.index_id, 5) def test_index_id_requires_table_id(self): with self.assertRaises(SystemExit): kss.parse_args(["10", "--index-id", "5"]) def test_conflicting_table_id_args(self): with self.assertRaises(SystemExit): kss.parse_args(["10", "123", "--table-id", "124"]) if __name__ == "__main__": unittest.main()