""" Tests for TiDB TopN value decoder This test module verifies the TiDB TopN value decoder can correctly decode various types of TiDB-encoded values including integers, floats, strings, etc. """ import struct import sys from utils.tidb_topn_decoder import ( BYTES_FLAG, COMPACT_BYTES_FLAG, FLOAT_FLAG, INT_FLAG, NIL_FLAG, UINT_FLAG, UVARINT_FLAG, VARINT_FLAG, decode_tidb_topn_value, decode_tidb_topn_value_safe, is_tidb_encoded_value, ) class TestTiDBTopNDecoder: """Test cases for TiDB TopN value decoder""" def validate_decode_signed_integer(self): """Test decoding signed integers""" # Test case from user: b'\x03\x80\x00\x00\x00\x00\x00\x00\x01' should decode to 1 encoded = b"\x03\x80\x00\x00\x00\x00\x00\x00\x01" result = decode_tidb_topn_value(encoded) assert result == 1, f"Expected 1, got {result}" # Test zero encoded_zero = b"\x03\x80\x00\x00\x00\x00\x00\x00\x00" result_zero = decode_tidb_topn_value(encoded_zero) assert result_zero == 0, f"Expected 0, got {result_zero}" # Test negative number: -1 encoded_neg = b"\x03\x7f\xff\xff\xff\xff\xff\xff\xff" result_neg = decode_tidb_topn_value(encoded_neg) assert result_neg == -1, f"Expected -1, got {result_neg}" def validate_decode_unsigned_integer(self): """Test decoding unsigned integers""" # Test small unsigned integer encoded = b"\x04\x00\x00\x00\x00\x00\x00\x00\x01" result = decode_tidb_topn_value(encoded) assert result == 1, f"Expected 1, got {result}" # Test larger unsigned integer encoded_large = b"\x04\x00\x00\x00\x00\x00\x00\x03\xe8" # 1000 result_large = decode_tidb_topn_value(encoded_large) assert result_large == 1000, f"Expected 1000, got {result_large}" def validate_decode_float(self): """Test decoding floating point numbers""" # Test positive float: 3.14 # First encode 3.14 as IEEE 754 double, then apply TiDB encoding original_float = 3.14 ieee_bytes = struct.pack(">d", original_float) # Apply TiDB encoding (flip sign bit for positive numbers) tidb_bytes = bytearray(ieee_bytes) tidb_bytes[0] |= 0x80 encoded = bytes([FLOAT_FLAG]) + bytes(tidb_bytes) result = decode_tidb_topn_value(encoded) assert abs(result - 3.14) < 0.001, f"Expected ~3.14, got {result}" def validate_decode_string_bytes(self): """Test decoding string values encoded as bytes""" # Test simple string - this is a simplified test case # Real TiDB encoding is more complex with group markers validate_string = "hello" # Create a simplified encoded bytes representation # In real TiDB, this would use group encoding, but for testing we'll use a simpler format string_bytes = validate_string.encode("utf-8") # Pad to 8 bytes and add marker padded = string_bytes.ljust(8, b"\x00") marker = 0xFF - (8 - len(string_bytes)) # Calculate marker based on padding encoded = bytes([BYTES_FLAG]) + padded + bytes([marker]) result = decode_tidb_topn_value(encoded) assert isinstance(result, str) assert validate_string in result or result.startswith(validate_string) def validate_decode_compact_bytes(self): """Test decoding compact byte strings with varint length""" validate_string = "world" string_bytes = validate_string.encode("utf-8") # Encode length as varint (simple case: length < 128) length = len(string_bytes) varint_length = bytes( [length] ) # For length < 128, varint is just the byte itself encoded = bytes([COMPACT_BYTES_FLAG]) + varint_length + string_bytes result = decode_tidb_topn_value(encoded) assert ( result == validate_string ), f"Expected '{validate_string}', got '{result}'" def validate_decode_nil_value(self): """Test decoding NULL values""" encoded = bytes([NIL_FLAG]) result = decode_tidb_topn_value(encoded) assert result is None, f"Expected None, got {result}" def validate_decode_varint(self): """Test decoding variable-length signed integers""" # Test small positive number (encoded as varint) # For varint: even values are positive/2, odd values are -(value+1)/2 # So 1 encoded as varint would be: 1 -> 2 -> 0x02 encoded = bytes([VARINT_FLAG, 0x02]) # Represents 1 result = decode_tidb_topn_value(encoded) assert result == 1, f"Expected 1, got {result}" # Test zero encoded_zero = bytes([VARINT_FLAG, 0x00]) result_zero = decode_tidb_topn_value(encoded_zero) assert result_zero == 0, f"Expected 0, got {result_zero}" def validate_decode_uvarint(self): """Test decoding variable-length unsigned integers""" # Test small unsigned number encoded = bytes([UVARINT_FLAG, 0x05]) # Represents 5 result = decode_tidb_topn_value(encoded) assert result == 5, f"Expected 5, got {result}" def validate_is_tidb_encoded_value(self): """Test detection of TiDB-encoded values""" # Valid TiDB encoded value valid_encoded = b"\x03\x80\x00\x00\x00\x00\x00\x00\x01" assert is_tidb_encoded_value(valid_encoded) # Invalid/random bytes invalid_bytes = b"\xff\x01\x02\x03" assert not is_tidb_encoded_value(invalid_bytes) # Empty bytes assert not is_tidb_encoded_value(b"") assert not is_tidb_encoded_value(None) def validate_decode_tidb_topn_value_safe(self): """Test safe decoding with error handling""" # Valid value valid_encoded = b"\x03\x80\x00\x00\x00\x00\x00\x00\x01" result = decode_tidb_topn_value_safe(valid_encoded) assert result == 1 # Invalid value - should return default invalid_encoded = b"\xff\x01\x02" result = decode_tidb_topn_value_safe(invalid_encoded, default="fallback") assert result == "fallback" # Empty value result_empty = decode_tidb_topn_value_safe(b"", default=None) assert result_empty is None def validate_edge_cases(self): """Test edge cases and error conditions""" # Empty encoded value try: decode_tidb_topn_value(b"") assert False, "Expected ValueError for empty encoded value" except ValueError as e: assert "Empty encoded value" in str(e) # Unknown type flag try: decode_tidb_topn_value(b"\xff\x01\x02\x03") assert False, "Expected ValueError for unsupported type flag" except ValueError as e: assert "Unsupported type flag" in str(e) # Insufficient data for integer try: decode_tidb_topn_value(b"\x03\x01\x02") # INT_FLAG but only 2 bytes of data assert False, "Expected ValueError for insufficient data" except ValueError as e: assert "Invalid signed int data length" in str(e) def validate_real_world_examples(self): """Test with examples that might appear in real TiDB data""" # Test the original problem case from user problematic_value = b"\x03\x80\x00\x00\x00\x00\x00\x00\x01" result = decode_tidb_topn_value(problematic_value) assert result == 1, f"Failed to decode original problematic value, got {result}" # Test various integer values that might appear in TopN validate_cases = [ (b"\x03\x80\x00\x00\x00\x00\x00\x00\x00", 0), # 0 (b"\x03\x80\x00\x00\x00\x00\x00\x00\x01", 1), # 1 (b"\x03\x80\x00\x00\x00\x00\x00\x00\x0a", 10), # 10 (b"\x03\x80\x00\x00\x00\x00\x00\x00\x64", 100), # 100 ] for encoded_bytes, expected_value in validate_cases: result = decode_tidb_topn_value(encoded_bytes) assert ( result == expected_value ), f"Expected {expected_value}, got {result} for {encoded_bytes.hex()}" def validate_chinese_text_handling(self): """测试中文文本处理 - Test Chinese text in test cases as allowed""" # This test case can contain Chinese text as per requirements validate_cases = [ "你好世界", # Hello World in Chinese "测试数据", # Test Data in Chinese "中文编码", # Chinese Encoding in Chinese ] for chinese_text in validate_cases: # Simulate how TiDB might encode Chinese text text_bytes = chinese_text.encode("utf-8") # Create compact bytes encoding length = len(text_bytes) if length < 128: varint_length = bytes([length]) else: # Handle longer strings with proper varint encoding varint_length = bytes([length & 0x7F | 0x80, length >> 7]) encoded = bytes([COMPACT_BYTES_FLAG]) + varint_length + text_bytes result = decode_tidb_topn_value(encoded) assert result == chinese_text, f"Expected '{chinese_text}', got '{result}'" def validate_type_flag_coverage(self): """Test coverage of all supported type flags""" # Test that all documented type flags are handled without raising "Unsupported type flag" supported_flags = [ NIL_FLAG, BYTES_FLAG, COMPACT_BYTES_FLAG, INT_FLAG, UINT_FLAG, FLOAT_FLAG, VARINT_FLAG, UVARINT_FLAG, ] for flag in supported_flags: try: # Create minimal valid data for each flag type if flag == NIL_FLAG: encoded = bytes([flag]) elif flag in [INT_FLAG, UINT_FLAG, FLOAT_FLAG]: encoded = bytes([flag]) + b"\x00" * 8 elif flag in [BYTES_FLAG]: encoded = bytes([flag]) + b"test\x00\x00\x00\x00\xff" elif flag == COMPACT_BYTES_FLAG: encoded = bytes([flag]) + b"\x04test" elif flag in [VARINT_FLAG, UVARINT_FLAG]: encoded = bytes([flag]) + b"\x00" # Should not raise "Unsupported type flag" error decode_tidb_topn_value(encoded) except ValueError as e: # Other ValueErrors are acceptable (like insufficient data) # but "Unsupported type flag" should not occur assert "Unsupported type flag" not in str( e ), f"Flag {flag} should be supported" except Exception: # Other exceptions are acceptable for this test # We're only testing that the type flag is recognized pass def run_all_tests(): """Run all test methods""" validate_instance = TestTiDBTopNDecoder() validate_methods = [ validate_instance.validate_decode_signed_integer, validate_instance.validate_decode_unsigned_integer, validate_instance.validate_decode_float, validate_instance.validate_decode_string_bytes, validate_instance.validate_decode_compact_bytes, validate_instance.validate_decode_nil_value, validate_instance.validate_decode_varint, validate_instance.validate_decode_uvarint, validate_instance.validate_is_tidb_encoded_value, validate_instance.validate_decode_tidb_topn_value_safe, validate_instance.validate_edge_cases, validate_instance.validate_real_world_examples, validate_instance.validate_chinese_text_handling, validate_instance.validate_type_flag_coverage, ] passed = 0 failed = 0 for validate_method in validate_methods: try: print(f"Running {validate_method.__name__}...") validate_method() print(f"✓ {validate_method.__name__} passed") passed += 1 except Exception as e: print(f"✗ {validate_method.__name__} failed: {e}") failed += 1 print(f"\nTest Results: {passed} passed, {failed} failed") return failed == 0 if __name__ == "__main__": success = run_all_tests() sys.exit(0 if success else 1)