diff --git a/tests/unit_tests/test_timelock.py b/tests/unit_tests/test_timelock.py new file mode 100644 index 0000000000..607bb4e429 --- /dev/null +++ b/tests/unit_tests/test_timelock.py @@ -0,0 +1,503 @@ +""" +Unit tests for bittensor.core.timelock module. + +This module provides comprehensive unit tests for the TimeLock Encryption (TLE) functionality, +which encrypts data such that it can only be decrypted after a specific amount of time +(expressed in Drand rounds). + +Test Coverage: +- encrypt() function: input validation, string/bytes handling, parameter passing +- decrypt() function: error handling, return types, no_errors flag +- wait_reveal_and_decrypt() function: round parsing, waiting logic, error handling +- TLE_ENCRYPTED_DATA_SUFFIX constant validation +- Edge cases and error conditions + +Contribution by Gittensor, learn more at https://gittensor.io/ +""" + +import struct +from unittest.mock import patch + +import pytest + +from bittensor.core import timelock +from bittensor.core.timelock import ( + encrypt, + decrypt, + wait_reveal_and_decrypt, + TLE_ENCRYPTED_DATA_SUFFIX, +) + + +class TestTLEConstants: + """Tests for module-level constants.""" + + def test_tle_encrypted_data_suffix_is_bytes(self): + """Test that TLE_ENCRYPTED_DATA_SUFFIX is bytes type.""" + assert isinstance(TLE_ENCRYPTED_DATA_SUFFIX, bytes) + + def test_tle_encrypted_data_suffix_value(self): + """Test that TLE_ENCRYPTED_DATA_SUFFIX has expected value.""" + assert TLE_ENCRYPTED_DATA_SUFFIX == b"AES_GCM_" + + def test_module_exports(self): + """Test that __all__ exports expected functions.""" + assert "encrypt" in timelock.__all__ + assert "decrypt" in timelock.__all__ + assert "wait_reveal_and_decrypt" in timelock.__all__ + assert "get_latest_round" in timelock.__all__ + + +class TestEncrypt: + """Tests for the encrypt() function.""" + + @patch("bittensor.core.timelock._btr_encrypt") + def test_encrypt_with_string_data(self, mock_btr_encrypt): + """Test that encrypt() converts string to bytes before encryption.""" + mock_btr_encrypt.return_value = (b"encrypted_data", 12345) + + result = encrypt("test string", n_blocks=5) + + # Verify string was encoded to bytes + mock_btr_encrypt.assert_called_once_with(b"test string", 5, 12.0) + assert result == (b"encrypted_data", 12345) + + @patch("bittensor.core.timelock._btr_encrypt") + def test_encrypt_with_bytes_data(self, mock_btr_encrypt): + """Test that encrypt() passes bytes data directly.""" + mock_btr_encrypt.return_value = (b"encrypted_data", 12345) + + result = encrypt(b"test bytes", n_blocks=5) + + mock_btr_encrypt.assert_called_once_with(b"test bytes", 5, 12.0) + assert result == (b"encrypted_data", 12345) + + @patch("bittensor.core.timelock._btr_encrypt") + def test_encrypt_with_default_block_time(self, mock_btr_encrypt): + """Test that encrypt() uses default block_time of 12.0 seconds.""" + mock_btr_encrypt.return_value = (b"encrypted", 100) + + encrypt(b"data", n_blocks=10) + + mock_btr_encrypt.assert_called_once_with(b"data", 10, 12.0) + + @patch("bittensor.core.timelock._btr_encrypt") + def test_encrypt_with_custom_block_time(self, mock_btr_encrypt): + """Test that encrypt() passes custom block_time correctly.""" + mock_btr_encrypt.return_value = (b"encrypted", 100) + + encrypt(b"data", n_blocks=10, block_time=0.25) + + mock_btr_encrypt.assert_called_once_with(b"data", 10, 0.25) + + @patch("bittensor.core.timelock._btr_encrypt") + def test_encrypt_with_fast_block_time(self, mock_btr_encrypt): + """Test encrypt() with fast-blocks mode (block_time=0.25).""" + mock_btr_encrypt.return_value = (b"fast_encrypted", 500) + + result = encrypt("fast mode", n_blocks=15, block_time=0.25) + + mock_btr_encrypt.assert_called_once_with(b"fast mode", 15, 0.25) + assert result == (b"fast_encrypted", 500) + + @patch("bittensor.core.timelock._btr_encrypt") + def test_encrypt_with_integer_block_time(self, mock_btr_encrypt): + """Test that encrypt() accepts integer block_time.""" + mock_btr_encrypt.return_value = (b"encrypted", 100) + + encrypt(b"data", n_blocks=5, block_time=6) + + mock_btr_encrypt.assert_called_once_with(b"data", 5, 6) + + @patch("bittensor.core.timelock._btr_encrypt") + def test_encrypt_returns_tuple(self, mock_btr_encrypt): + """Test that encrypt() returns a tuple of (bytes, int).""" + mock_btr_encrypt.return_value = (b"encrypted_data", 99999) + + result = encrypt(b"test", n_blocks=1) + + assert isinstance(result, tuple) + assert len(result) == 2 + assert isinstance(result[0], bytes) + assert isinstance(result[1], int) + + @patch("bittensor.core.timelock._btr_encrypt") + def test_encrypt_with_empty_string(self, mock_btr_encrypt): + """Test encrypt() with empty string input.""" + mock_btr_encrypt.return_value = (b"encrypted_empty", 100) + + result = encrypt("", n_blocks=1) + + mock_btr_encrypt.assert_called_once_with(b"", 1, 12.0) + assert result == (b"encrypted_empty", 100) + + @patch("bittensor.core.timelock._btr_encrypt") + def test_encrypt_with_empty_bytes(self, mock_btr_encrypt): + """Test encrypt() with empty bytes input.""" + mock_btr_encrypt.return_value = (b"encrypted_empty", 100) + + encrypt(b"", n_blocks=1) + + mock_btr_encrypt.assert_called_once_with(b"", 1, 12.0) + + @patch("bittensor.core.timelock._btr_encrypt") + def test_encrypt_with_unicode_string(self, mock_btr_encrypt): + """Test encrypt() with unicode string input.""" + mock_btr_encrypt.return_value = (b"encrypted_unicode", 100) + unicode_str = "Hello δΈ–η•Œ 🌍" + + encrypt(unicode_str, n_blocks=1) + + mock_btr_encrypt.assert_called_once_with(unicode_str.encode(), 1, 12.0) + + @patch("bittensor.core.timelock._btr_encrypt") + def test_encrypt_with_large_n_blocks(self, mock_btr_encrypt): + """Test encrypt() with large n_blocks value.""" + mock_btr_encrypt.return_value = (b"encrypted", 1000000) + + encrypt(b"data", n_blocks=100000) + + mock_btr_encrypt.assert_called_once_with(b"data", 100000, 12.0) + + +class TestDecrypt: + """Tests for the decrypt() function.""" + + @patch("bittensor.core.timelock._btr_decrypt") + def test_decrypt_returns_bytes(self, mock_btr_decrypt): + """Test that decrypt() returns bytes by default.""" + mock_btr_decrypt.return_value = b"decrypted_data" + + result = decrypt(b"encrypted_data") + + mock_btr_decrypt.assert_called_once_with(b"encrypted_data", True) + assert result == b"decrypted_data" + assert isinstance(result, bytes) + + @patch("bittensor.core.timelock._btr_decrypt") + def test_decrypt_returns_none_before_reveal(self, mock_btr_decrypt): + """Test that decrypt() returns None when reveal round not reached.""" + mock_btr_decrypt.return_value = None + + result = decrypt(b"encrypted_data") + + assert result is None + + @patch("bittensor.core.timelock._btr_decrypt") + def test_decrypt_with_no_errors_true(self, mock_btr_decrypt): + """Test decrypt() with no_errors=True (default).""" + mock_btr_decrypt.return_value = b"data" + + decrypt(b"encrypted", no_errors=True) + + mock_btr_decrypt.assert_called_once_with(b"encrypted", True) + + @patch("bittensor.core.timelock._btr_decrypt") + def test_decrypt_with_no_errors_false(self, mock_btr_decrypt): + """Test decrypt() with no_errors=False.""" + mock_btr_decrypt.return_value = b"data" + + decrypt(b"encrypted", no_errors=False) + + mock_btr_decrypt.assert_called_once_with(b"encrypted", False) + + @patch("bittensor.core.timelock._btr_decrypt") + def test_decrypt_with_return_str_true(self, mock_btr_decrypt): + """Test decrypt() with return_str=True returns string.""" + mock_btr_decrypt.return_value = b"decrypted_string" + + result = decrypt(b"encrypted", return_str=True) + + assert result == "decrypted_string" + assert isinstance(result, str) + + @patch("bittensor.core.timelock._btr_decrypt") + def test_decrypt_with_return_str_false(self, mock_btr_decrypt): + """Test decrypt() with return_str=False returns bytes.""" + mock_btr_decrypt.return_value = b"decrypted_bytes" + + result = decrypt(b"encrypted", return_str=False) + + assert result == b"decrypted_bytes" + assert isinstance(result, bytes) + + @patch("bittensor.core.timelock._btr_decrypt") + def test_decrypt_return_str_with_none_result(self, mock_btr_decrypt): + """Test decrypt() with return_str=True when result is None.""" + mock_btr_decrypt.return_value = None + + result = decrypt(b"encrypted", return_str=True) + + assert result is None + + @patch("bittensor.core.timelock._btr_decrypt") + def test_decrypt_with_unicode_result(self, mock_btr_decrypt): + """Test decrypt() with unicode bytes result and return_str=True.""" + mock_btr_decrypt.return_value = "Hello δΈ–η•Œ".encode() + + result = decrypt(b"encrypted", return_str=True) + + assert result == "Hello δΈ–η•Œ" + + +class TestWaitRevealAndDecrypt: + """Tests for the wait_reveal_and_decrypt() function.""" + + @patch("bittensor.core.timelock.decrypt") + @patch("bittensor.core.timelock.get_latest_round") + def test_wait_reveal_with_explicit_round(self, mock_get_round, mock_decrypt): + """Test wait_reveal_and_decrypt() with explicit reveal_round.""" + mock_get_round.return_value = 100 # Already past reveal round + mock_decrypt.return_value = b"decrypted" + + result = wait_reveal_and_decrypt(b"encrypted", reveal_round=50) + + mock_decrypt.assert_called_once_with(b"encrypted", True, False) + assert result == b"decrypted" + + @patch("bittensor.core.timelock.decrypt") + @patch("bittensor.core.timelock.get_latest_round") + def test_wait_reveal_parses_round_from_data(self, mock_get_round, mock_decrypt): + """Test wait_reveal_and_decrypt() parses reveal_round from encrypted data.""" + # Create mock encrypted data with embedded round + reveal_round = 12345 + encrypted_data = ( + b"some_encrypted_data" + + TLE_ENCRYPTED_DATA_SUFFIX + + struct.pack("