brainstate 0.1.10__py2.py3-none-any.whl → 0.2.1__py2.py3-none-any.whl

brainstate/environ_test.py

@@ -1,62 +1,1223 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-
-import unittest
-
-import jax.numpy as jnp
-
-import brainstate as bst
-
-
-class TestEnviron(unittest.TestCase):
-    def test_precision(self):
-        with bst.environ.context(precision=64):
-            a = bst.random.randn(1)
-            self.assertEqual(a.dtype, jnp.float64)
-
-        with bst.environ.context(precision=32):
-            a = bst.random.randn(1)
-            self.assertEqual(a.dtype, jnp.float32)
-
-        with bst.environ.context(precision=16):
-            a = bst.random.randn(1)
-            self.assertEqual(a.dtype, jnp.float16)
-
-        with bst.environ.context(precision='bf16'):
-            a = bst.random.randn(1)
-            self.assertEqual(a.dtype, jnp.bfloat16)
-
-    def test_platform(self):
-        with self.assertRaises(ValueError):
-            with bst.environ.context(platform='cpu'):
-                a = bst.random.randn(1)
-                self.assertEqual(a.device(), 'cpu')
-
-    def test_register_default_behavior(self):
-        bst.environ.set(dt=0.1)
-
-        dt_ = 0.1
-
-        def dt_behavior(dt):
-            nonlocal dt_
-            dt_ = dt
-            print(f'dt: {dt}')
-
-        bst.environ.register_default_behavior('dt', dt_behavior)
-
-        with bst.environ.context(dt=0.2):
-            self.assertEqual(dt_, 0.2)
-        self.assertEqual(dt_, 0.1)
+# Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""
+Comprehensive test suite for the environ module.
+
+This test module provides extensive coverage of the environment configuration
+and context management functionality, including:
+- Global environment settings
+- Context-based temporary settings
+- Precision and data type management
+- Callback registration and behavior
+- Thread safety
+- Error handling and validation
+"""
+
+import threading
+import unittest
+import warnings
+from unittest.mock import patch
+
+import jax.numpy as jnp
+import numpy as np
+
+import brainstate as bst
+
+
+class TestEnvironmentCore(unittest.TestCase):
+    """Test core environment management functionality."""
+
+    def setUp(self):
+        """Reset environment before each test."""
+        bst.environ.reset()
+        # Clear any warnings
+        warnings.filterwarnings('ignore', category=UserWarning)
+
+    def tearDown(self):
+        """Clean up after each test."""
+        # Reset to default state
+        bst.environ.reset()
+        warnings.resetwarnings()
+
+    def test_get_set_basic(self):
+        """Test basic get and set operations."""
+        # Set a value
+        bst.environ.set(test_param='test_value')
+        self.assertEqual(bst.environ.get('test_param'), 'test_value')
+
+        # Set multiple values
+        bst.environ.set(param1=1, param2='two', param3=3.0)
+        self.assertEqual(bst.environ.get('param1'), 1)
+        self.assertEqual(bst.environ.get('param2'), 'two')
+        self.assertEqual(bst.environ.get('param3'), 3.0)
+
+    def test_get_with_default(self):
+        """Test get with default value."""
+        # Non-existent key with default
+        result = bst.environ.get('nonexistent', default='default_value')
+        self.assertEqual(result, 'default_value')
+
+        # Existing key ignores default
+        bst.environ.set(existing='value')
+        result = bst.environ.get('existing', default='default')
+        self.assertEqual(result, 'value')
+
+    def test_get_missing_key_error(self):
+        """Test KeyError for missing keys without default."""
+        with self.assertRaises(KeyError) as context:
+            bst.environ.get('missing_key')
+
+        error_msg = str(context.exception)
+        self.assertIn('missing_key', error_msg)
+        self.assertIn('not found', error_msg)
+
+    def test_get_with_description(self):
+        """Test get with description for error messages."""
+        with self.assertRaises(KeyError) as context:
+            bst.environ.get('missing', desc='Important parameter for computation')
+
+        error_msg = str(context.exception)
+        self.assertIn('Important parameter', error_msg)
+
+    def test_all_function(self):
+        """Test getting all environment settings."""
+        # Set various parameters
+        bst.environ.set(
+            param1='value1',
+            param2=42,
+            param3=3.14,
+            precision=32
+        )
+
+        all_settings = bst.environ.all()
+        self.assertIsInstance(all_settings, dict)
+        self.assertEqual(all_settings['param1'], 'value1')
+        self.assertEqual(all_settings['param2'], 42)
+        self.assertEqual(all_settings['param3'], 3.14)
+        self.assertEqual(all_settings['precision'], 32)
+
+    def test_reset_function(self):
+        """Test environment reset functionality."""
+        # Set custom values
+        bst.environ.set(
+            custom1='value1',
+            custom2='value2',
+            precision=64
+        )
+
+        # Verify they're set
+        self.assertEqual(bst.environ.get('custom1'), 'value1')
+
+        # Reset environment
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            bst.environ.reset()
+
+        # Custom values should be gone
+        result = bst.environ.get('custom1', default=None)
+        self.assertIsNone(result)
+
+        # Default precision should be restored
+        self.assertEqual(bst.environ.get('precision'), bst.environ.DEFAULT_PRECISION)
+
+    def test_special_environment_keys(self):
+        """Test special environment key constants."""
+        # Test setting using constants
+        bst.environ.set(**{
+            bst.environ.DT: 0.01,
+            bst.environ.I: 0,
+            bst.environ.T: 0.0,
+            bst.environ.JIT_ERROR_CHECK: True,
+            bst.environ.FIT: False
+        })
+
+        self.assertEqual(bst.environ.get(bst.environ.DT), 0.01)
+        self.assertEqual(bst.environ.get(bst.environ.I), 0)
+        self.assertEqual(bst.environ.get(bst.environ.T), 0.0)
+        self.assertTrue(bst.environ.get(bst.environ.JIT_ERROR_CHECK))
+        self.assertFalse(bst.environ.get(bst.environ.FIT))
+
+    def test_pop_basic(self):
+        """Test basic pop operation."""
+        # Set a value
+        bst.environ.set(pop_test='test_value')
+        self.assertEqual(bst.environ.get('pop_test'), 'test_value')
+
+        # Pop the value
+        popped = bst.environ.pop('pop_test')
+        self.assertEqual(popped, 'test_value')
+
+        # Value should be gone
+        result = bst.environ.get('pop_test', default=None)
+        self.assertIsNone(result)
+
+    def test_pop_with_default(self):
+        """Test pop with default value."""
+        # Pop non-existent key with default
+        result = bst.environ.pop('nonexistent_pop', default='default_value')
+        self.assertEqual(result, 'default_value')
+
+        # Pop existing key ignores default
+        bst.environ.set(existing_pop='value')
+        result = bst.environ.pop('existing_pop', default='default')
+        self.assertEqual(result, 'value')
+
+    def test_pop_missing_key_error(self):
+        """Test KeyError for missing keys without default."""
+        with self.assertRaises(KeyError) as context:
+            bst.environ.pop('missing_pop_key')
+
+        error_msg = str(context.exception)
+        self.assertIn('missing_pop_key', error_msg)
+        self.assertIn('not found', error_msg)
+
+    def test_pop_multiple_values(self):
+        """Test popping multiple values."""
+        # Set multiple values
+        bst.environ.set(
+            pop1='value1',
+            pop2='value2',
+            pop3='value3'
+        )
+
+        # Pop them one by one
+        v1 = bst.environ.pop('pop1')
+        v2 = bst.environ.pop('pop2')
+
+        self.assertEqual(v1, 'value1')
+        self.assertEqual(v2, 'value2')
+
+        # pop3 should still exist
+        self.assertEqual(bst.environ.get('pop3'), 'value3')
+
+        # pop1 and pop2 should be gone
+        self.assertIsNone(bst.environ.get('pop1', default=None))
+        self.assertIsNone(bst.environ.get('pop2', default=None))
+
+    def test_pop_with_context_protection(self):
+        """Test that pop is prevented when key is in active context."""
+        # Set a global value
+        bst.environ.set(protected_key='global_value')
+
+        # Cannot pop while in context
+        with bst.environ.context(protected_key='context_value'):
+            with self.assertRaises(ValueError) as context:
+                bst.environ.pop('protected_key')
+
+            error_msg = str(context.exception)
+            self.assertIn('Cannot pop', error_msg)
+            self.assertIn('active in a context', error_msg)
+
+        # Can pop after context exits
+        popped = bst.environ.pop('protected_key')
+        self.assertEqual(popped, 'global_value')
+
+    def test_pop_nested_context_protection(self):
+        """Test pop protection with nested contexts."""
+        bst.environ.set(nested_key='global')
+
+        with bst.environ.context(nested_key='level1'):
+            with bst.environ.context(nested_key='level2'):
+                # Should indicate 2 active contexts
+                with self.assertRaises(ValueError) as context:
+                    bst.environ.pop('nested_key')
+
+                error_msg = str(context.exception)
+                self.assertIn('2 context(s)', error_msg)
+
+    def test_pop_does_not_affect_context_values(self):
+        """Test that pop doesn't affect context values."""
+        # Set both global and context value
+        bst.environ.set(dual_key='global')
+
+        with bst.environ.context(other_key='context_only'):
+            # Can pop a key that's only in global (not in this context)
+            popped = bst.environ.pop('dual_key')
+            self.assertEqual(popped, 'global')
+
+            # Context-only values remain accessible
+            self.assertEqual(bst.environ.get('other_key'), 'context_only')
+
+        # Context value should be gone after exit
+        self.assertIsNone(bst.environ.get('other_key', default=None))
+
+    def test_pop_precision_key(self):
+        """Test popping the precision key."""
+        # Set custom precision
+        bst.environ.set(precision=64)
+        self.assertEqual(bst.environ.get_precision(), 64)
+
+        # Pop precision
+        popped = bst.environ.pop('precision')
+        self.assertEqual(popped, 64)
+
+
+class TestEnvironmentContext(unittest.TestCase):
+    """Test context manager functionality."""
+
+    def setUp(self):
+        """Reset environment before each test."""
+        bst.environ.reset()
+        warnings.filterwarnings('ignore', category=UserWarning)
+
+    def tearDown(self):
+        """Clean up after each test."""
+        bst.environ.reset()
+        warnings.resetwarnings()
+
+    def test_basic_context(self):
+        """Test basic context manager usage."""
+        bst.environ.set(value=10)
+
+        with bst.environ.context(value=20) as ctx:
+            # Value should be 20 in context
+            self.assertEqual(bst.environ.get('value'), 20)
+            # Context should contain current settings
+            self.assertEqual(ctx['value'], 20)
+
+        # Value should be restored to 10
+        self.assertEqual(bst.environ.get('value'), 10)
+
+    def test_nested_contexts(self):
+        """Test nested context managers."""
+        bst.environ.set(level=0)
+
+        with bst.environ.context(level=1):
+            self.assertEqual(bst.environ.get('level'), 1)
+
+            with bst.environ.context(level=2):
+                self.assertEqual(bst.environ.get('level'), 2)
+
+                with bst.environ.context(level=3):
+                    self.assertEqual(bst.environ.get('level'), 3)
+
+                # Back to level 2
+                self.assertEqual(bst.environ.get('level'), 2)
+
+            # Back to level 1
+            self.assertEqual(bst.environ.get('level'), 1)
+
+        # Back to level 0
+        self.assertEqual(bst.environ.get('level'), 0)
+
+    def test_context_with_exception(self):
+        """Test context manager handles exceptions properly."""
+        bst.environ.set(value='original')
+
+        try:
+            with bst.environ.context(value='temporary'):
+                self.assertEqual(bst.environ.get('value'), 'temporary')
+                raise ValueError("Test exception")
+        except ValueError:
+            pass
+
+        # Value should be restored despite exception
+        self.assertEqual(bst.environ.get('value'), 'original')
+
+    def test_context_multiple_parameters(self):
+        """Test context with multiple parameters."""
+        bst.environ.set(a=1, b=2, c=3)
+
+        with bst.environ.context(a=10, b=20, c=30, d=40):
+            self.assertEqual(bst.environ.get('a'), 10)
+            self.assertEqual(bst.environ.get('b'), 20)
+            self.assertEqual(bst.environ.get('c'), 30)
+            self.assertEqual(bst.environ.get('d'), 40)
+
+        # Original values restored
+        self.assertEqual(bst.environ.get('a'), 1)
+        self.assertEqual(bst.environ.get('b'), 2)
+        self.assertEqual(bst.environ.get('c'), 3)
+        # d should not exist
+        result = bst.environ.get('d', default=None)
+        self.assertIsNone(result)
+
+    def test_context_platform_restriction(self):
+        """Test that platform cannot be set in context."""
+        with self.assertRaises(ValueError) as context:
+            with bst.environ.context(platform='cpu'):
+                pass
+
+        self.assertIn('platform', str(context.exception).lower())
+        self.assertIn('cannot set', str(context.exception).lower())
+
+    def test_context_host_device_count_restriction(self):
+        """Test that host_device_count cannot be set in context."""
+        with self.assertRaises(ValueError) as context:
+            with bst.environ.context(host_device_count=4):
+                pass
+
+        self.assertIn('host_device_count', str(context.exception))
+
+    def test_context_mode_validation(self):
+        """Test mode validation in context."""
+        # Valid mode
+        mode = bst.mixin.Training()
+        with bst.environ.context(mode=mode):
+            self.assertEqual(bst.environ.get('mode'), mode)
+
+    def test_context_preserves_unmodified_values(self):
+        """Test that context doesn't affect unmodified values."""
+        bst.environ.set(unchanged='original', changed='original')
+
+        with bst.environ.context(changed='modified'):
+            self.assertEqual(bst.environ.get('unchanged'), 'original')
+            self.assertEqual(bst.environ.get('changed'), 'modified')
+
+
+class TestPrecisionAndDataTypes(unittest.TestCase):
+    """Test precision control and data type functions."""
+
+    def setUp(self):
+        """Reset environment before each test."""
+        bst.environ.reset()
+        warnings.filterwarnings('ignore', category=UserWarning)
+
+    def tearDown(self):
+        """Clean up after each test."""
+        bst.environ.reset()
+        warnings.resetwarnings()
+
+    def test_precision_settings(self):
+        """Test different precision settings."""
+        precisions = [8, 16, 32, 64, 'bf16']
+
+        for precision in precisions:
+            bst.environ.set(precision=precision)
+
+            if precision == 'bf16':
+                self.assertEqual(bst.environ.get_precision(), 16)
+            elif isinstance(precision, str):
+                self.assertEqual(bst.environ.get_precision(), int(precision))
+            else:
+                self.assertEqual(bst.environ.get_precision(), precision)
+
+    def test_precision_context(self):
+        """Test precision changes in context."""
+        bst.environ.set(precision=32)
+
+        with bst.environ.context(precision=64):
+            a = bst.random.randn(1)
+            self.assertEqual(a.dtype, jnp.float64)
+            self.assertEqual(bst.environ.get_precision(), 64)
+
+        # Precision restored
+        b = bst.random.randn(1)
+        self.assertEqual(b.dtype, jnp.float32)
+        self.assertEqual(bst.environ.get_precision(), 32)
+
+    def test_dftype_function(self):
+        """Test default float type function."""
+        # 32-bit precision
+        bst.environ.set(precision=32)
+        self.assertEqual(bst.environ.dftype(), np.float32)
+
+        # 64-bit precision
+        bst.environ.set(precision=64)
+        self.assertEqual(bst.environ.dftype(), np.float64)
+
+        # 16-bit precision
+        bst.environ.set(precision=16)
+        self.assertEqual(bst.environ.dftype(), np.float16)
+
+        # bfloat16 precision
+        bst.environ.set(precision='bf16')
+        self.assertEqual(bst.environ.dftype(), jnp.bfloat16)
+
+    def test_ditype_function(self):
+        """Test default integer type function."""
+        # 32-bit precision
+        bst.environ.set(precision=32)
+        self.assertEqual(bst.environ.ditype(), np.int32)
+
+        # 64-bit precision
+        bst.environ.set(precision=64)
+        self.assertEqual(bst.environ.ditype(), np.int64)
+
+        # 16-bit precision
+        bst.environ.set(precision=16)
+        self.assertEqual(bst.environ.ditype(), np.int16)
+
+        # 8-bit precision
+        bst.environ.set(precision=8)
+        self.assertEqual(bst.environ.ditype(), np.int8)
+
+    def test_dutype_function(self):
+        """Test default unsigned integer type function."""
+        # 32-bit precision
+        bst.environ.set(precision=32)
+        self.assertEqual(bst.environ.dutype(), np.uint32)
+
+        # 64-bit precision
+        bst.environ.set(precision=64)
+        self.assertEqual(bst.environ.dutype(), np.uint64)
+
+        # 16-bit precision
+        bst.environ.set(precision=16)
+        self.assertEqual(bst.environ.dutype(), np.uint16)
+
+        # 8-bit precision
+        bst.environ.set(precision=8)
+        self.assertEqual(bst.environ.dutype(), np.uint8)
+
+    def test_dctype_function(self):
+        """Test default complex type function."""
+        # 32-bit precision
+        bst.environ.set(precision=32)
+        self.assertEqual(bst.environ.dctype(), np.complex64)
+
+        # 64-bit precision
+        bst.environ.set(precision=64)
+        self.assertEqual(bst.environ.dctype(), np.complex128)
+
+        # 16-bit precision (should use complex64)
+        bst.environ.set(precision=16)
+        self.assertEqual(bst.environ.dctype(), np.complex64)
+
+    def test_tolerance_function(self):
+        """Test tolerance values for different precisions."""
+        # 64-bit precision
+        bst.environ.set(precision=64)
+        tol = bst.environ.tolerance()
+        self.assertAlmostEqual(float(tol), 1e-12, places=14)
+
+        # 32-bit precision
+        bst.environ.set(precision=32)
+        tol = bst.environ.tolerance()
+        self.assertAlmostEqual(float(tol), 1e-5, places=7)
+
+        # 16-bit precision
+        bst.environ.set(precision=16)
+        tol = bst.environ.tolerance()
+        self.assertAlmostEqual(float(tol), 1e-2, places=4)
+
+    def test_invalid_precision(self):
+        """Test invalid precision values."""
+        invalid_precisions = [128, 'invalid', -1, 3.14]
+
+        for invalid in invalid_precisions:
+            with self.assertRaises(ValueError):
+                bst.environ.set(precision=invalid)
+
+    def test_precision_with_arrays(self):
+        """Test that precision affects array creation."""
+        # Test with different precisions
+        test_cases = [
+            (32, jnp.float32),
+            (64, jnp.float64),
+            (16, jnp.float16),
+            ('bf16', jnp.bfloat16),
+        ]
+
+        for precision, expected_dtype in test_cases:
+            with bst.environ.context(precision=precision):
+                # Create array using random
+                arr = bst.random.randn(5)
+                self.assertEqual(arr.dtype, expected_dtype)
+
+
+class TestModeAndSpecialGetters(unittest.TestCase):
+    """Test mode management and special getter functions."""
+
+    def setUp(self):
+        """Reset environment before each test."""
+        bst.environ.reset()
+        warnings.filterwarnings('ignore', category=UserWarning)
+
+    def tearDown(self):
+        """Clean up after each test."""
+        bst.environ.reset()
+        warnings.resetwarnings()
+
+    def test_get_dt(self):
+        """Test get_dt function."""
+        # Set dt
+        bst.environ.set(dt=0.01)
+        self.assertEqual(bst.environ.get_dt(), 0.01)
+
+        # Test in context
+        with bst.environ.context(dt=0.001):
+            self.assertEqual(bst.environ.get_dt(), 0.001)
+
+        self.assertEqual(bst.environ.get_dt(), 0.01)
+
+        # Test missing dt
+        bst.environ.reset()
+        with self.assertRaises(KeyError):
+            bst.environ.get_dt()
+
+    def test_get_mode(self):
+        """Test get_mode function."""
+        # Set training mode
+        training = bst.mixin.Training()
+        bst.environ.set(mode=training)
+        mode = bst.environ.get('mode')
+        self.assertEqual(mode, training)
+        self.assertTrue(mode.has(bst.mixin.Training))
+
+        # Test with batching mode
+        batching = bst.mixin.Batching(batch_size=32)
+        with bst.environ.context(mode=batching):
+            mode = bst.environ.get('mode')
+            self.assertEqual(mode, batching)
+            self.assertTrue(mode.has(bst.mixin.Batching))
+            self.assertEqual(mode.batch_size, 32)
+
+        # Test missing mode
+        bst.environ.reset()
+        with self.assertRaises(KeyError):
+            bst.environ.get('mode')
+
+    def test_get_platform(self):
+        """Test get_platform function."""
+        platform = bst.environ.get_platform()
+        self.assertIn(platform, bst.environ.SUPPORTED_PLATFORMS)
+
+    def test_get_host_device_count(self):
+        """Test get_host_device_count function."""
+        count = bst.environ.get_host_device_count()
+        self.assertIsInstance(count, int)
+        self.assertGreaterEqual(count, 1)
+
+    def test_dt_validation(self):
+        """Test dt validation in set function."""
+        # Valid dt values
+        valid_dts = [0.01, 0.001, 1.0, 0.1]
+        for dt in valid_dts:
+            bst.environ.set(dt=dt)
+            self.assertEqual(bst.environ.get_dt(), dt)
+
+
+class TestPlatformAndDevice(unittest.TestCase):
+    """Test platform and device management."""
+
+    def setUp(self):
+        """Reset environment before each test."""
+        bst.environ.reset()
+        warnings.filterwarnings('ignore', category=UserWarning)
+
+    def tearDown(self):
+        """Clean up after each test."""
+        bst.environ.reset()
+        warnings.resetwarnings()
+
+    @patch('brainstate.environ.config')
+    def test_set_platform(self, mock_config):
+        """Test platform setting."""
+        platforms = ['cpu', 'gpu', 'tpu']
+
+        for platform in platforms:
+            bst.environ.set_platform(platform)
+            mock_config.update.assert_called_with("jax_platform_name", platform)
+
+        # Test invalid platform
+        with self.assertRaises(ValueError):
+            bst.environ.set_platform('invalid')
+
+    def test_set_platform_through_set(self):
+        """Test setting platform through general set function."""
+        with patch('brainstate.environ.config') as mock_config:
+            bst.environ.set(platform='gpu')
+            mock_config.update.assert_called_with("jax_platform_name", 'gpu')
+
+    def test_set_host_device_count(self):
+        """Test host device count setting."""
+        import os
+
+        # Set device count
+        bst.environ.set_host_device_count(4)
+        xla_flags = os.environ.get("XLA_FLAGS", "")
+        self.assertIn("--xla_force_host_platform_device_count=4", xla_flags)
+
+        # Update device count
+        bst.environ.set_host_device_count(8)
+        xla_flags = os.environ.get("XLA_FLAGS", "")
+        self.assertIn("--xla_force_host_platform_device_count=8", xla_flags)
+        self.assertNotIn("--xla_force_host_platform_device_count=4", xla_flags)
+
+        # Invalid device count
+        with self.assertRaises(ValueError):
+            bst.environ.set_host_device_count(0)
+
+        with self.assertRaises(ValueError):
+            bst.environ.set_host_device_count(-1)
+
+    def test_platform_context_restriction(self):
+        """Test that platform cannot be changed in context."""
+        with self.assertRaises(ValueError):
+            with bst.environ.context(platform='cpu'):
+                pass
+
+
+class TestCallbackBehavior(unittest.TestCase):
+    """Test callback registration and behavior."""
+
+    def setUp(self):
+        """Reset environment before each test."""
+        bst.environ.reset()
+        warnings.filterwarnings('ignore', category=UserWarning)
+        self.callback_values = []
+
+    def tearDown(self):
+        """Clean up after each test."""
+        bst.environ.reset()
+        warnings.resetwarnings()
+
+    # def test_register_callback(self):
+    #     """Test basic callback registration."""
+    #     def callback(value):
+    #         self.callback_values.append(value)
+    #
+    #     brainstate.environ.register_default_behavior('test_param', callback)
+    #
+    #     # Callback should be triggered on set
+    #     brainstate.environ.set(test_param='value1')
+    #     self.assertEqual(self.callback_values, ['value1'])
+    #
+    #     # Callback should be triggered on context enter/exit
+    #     with brainstate.environ.context(test_param='value2'):
+    #         self.assertEqual(self.callback_values, ['value1', 'value2'])
+    #
+    #     # Should restore previous value
+    #     self.assertEqual(self.callback_values, ['value1', 'value2', 'value1'])
+
+    def test_register_multiple_callbacks(self):
+        """Test registering callbacks for different keys."""
+        values_a = []
+        values_b = []
+
+        def callback_a(value):
+            values_a.append(value)
+
+        def callback_b(value):
+            values_b.append(value)
+
+        bst.environ.register_default_behavior('param_a', callback_a)
+        bst.environ.register_default_behavior('param_b', callback_b)
+
+        bst.environ.set(param_a='a1', param_b='b1')
+        self.assertEqual(values_a, ['a1'])
+        self.assertEqual(values_b, ['b1'])
+
+    def test_replace_callback(self):
+        """Test replacing existing callbacks."""
+
+        def callback1(value):
+            self.callback_values.append(f'cb1:{value}')
+
+        def callback2(value):
+            self.callback_values.append(f'cb2:{value}')
+
+        # Register first callback
+        bst.environ.register_default_behavior('param', callback1)
+
+        # Try to register second without replace flag
+        with self.assertRaises(ValueError):
+            bst.environ.register_default_behavior('param', callback2)
+
+        # Register with replace flag
+        bst.environ.register_default_behavior('param', callback2, replace_if_exist=True)
+
+        # Only second callback should be called
+        bst.environ.set(param='test')
+        self.assertEqual(self.callback_values, ['cb2:test'])
+
+    def test_unregister_callback(self):
+        """Test unregistering callbacks."""
+
+        def callback(value):
+            self.callback_values.append(value)
+
+        # Register and test
+        bst.environ.register_default_behavior('param', callback)
+        bst.environ.set(param='value1')
+        self.assertEqual(len(self.callback_values), 1)
+
+        # Unregister
+        removed = bst.environ.unregister_default_behavior('param')
+        self.assertTrue(removed)
+
+        # Callback should not be triggered
+        bst.environ.set(param='value2')
+        self.assertEqual(len(self.callback_values), 1)  # Still just one
+
+        # Unregister non-existent
+        removed = bst.environ.unregister_default_behavior('nonexistent')
+        self.assertFalse(removed)
+
+    def test_list_registered_behaviors(self):
+        """Test listing registered behaviors."""
+        # Initially empty or with system defaults
+        initial = bst.environ.list_registered_behaviors()
+
+        # Register some behaviors
+        bst.environ.register_default_behavior('param1', lambda x: None)
+        bst.environ.register_default_behavior('param2', lambda x: None)
+        bst.environ.register_default_behavior('param3', lambda x: None)
+
+        behaviors = bst.environ.list_registered_behaviors()
+        for param in ['param1', 'param2', 'param3']:
+            self.assertIn(param, behaviors)
+
+    def test_callback_exception_handling(self):
+        """Test that exceptions in callbacks are handled gracefully."""
+
+        def failing_callback(value):
+            raise RuntimeError(f"Intentional error: {value}")
+
+        bst.environ.register_default_behavior('param', failing_callback)
+
+        # Should not crash, but should warn
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            bst.environ.set(param='test')
+
+            # Should have a warning
+            self.assertTrue(len(w) > 0)
+            self.assertIn('Callback', str(w[0].message))
+            self.assertIn('exception', str(w[0].message))
+
+    def test_callback_validation(self):
+        """Test callback validation."""
+        # Non-callable
+        with self.assertRaises(TypeError):
+            bst.environ.register_default_behavior('param', 'not_callable')
+
+        # Non-string key
+        with self.assertRaises(TypeError):
+            bst.environ.register_default_behavior(123, lambda x: None)
+
+    def test_callback_with_validation(self):
+        """Test using callbacks for validation."""
+
+        def validate_positive(value):
+            if value <= 0:
+                raise ValueError(f"Value must be positive, got {value}")
+            self.callback_values.append(value)
+
+        bst.environ.register_default_behavior('positive_param', validate_positive)
+
+        # Valid value
+        bst.environ.set(positive_param=10)
+        self.assertEqual(self.callback_values, [10])
+
+        # Invalid value should raise through warning system
+        with warnings.catch_warnings(record=True):
+            warnings.simplefilter("always")
+            bst.environ.set(positive_param=-5)
+
+
+class TestThreadSafety(unittest.TestCase):
+    """Test thread safety of environment operations."""
+
+    def setUp(self):
+        """Reset environment before each test."""
+        bst.environ.reset()
+        warnings.filterwarnings('ignore', category=UserWarning)
+
+    def tearDown(self):
+        """Clean up after each test."""
+        bst.environ.reset()
+        warnings.resetwarnings()
+
+    def test_concurrent_set_operations(self):
+        """Test concurrent set operations from multiple threads."""
+        results = []
+        errors = []
+
+        def thread_operation(thread_id):
+            try:
+                # Each thread sets its own value
+                for i in range(10):
+                    bst.environ.set(**{f'thread_{thread_id}': i})
+                    value = bst.environ.get(f'thread_{thread_id}')
+                    results.append((thread_id, value))
+            except Exception as e:
+                errors.append(e)
+
+        threads = []
+        for i in range(5):
+            thread = threading.Thread(target=thread_operation, args=(i,))
+            threads.append(thread)
+            thread.start()
+
+        for thread in threads:
+            thread.join()
+
+        # Should have no errors
+        self.assertEqual(len(errors), 0)
+
+        # Each thread should have written its values
+        for i in range(5):
+            try:
+                final_value = bst.environ.get(f'thread_{i}')
+            except KeyError:
+                pass
+
+    def test_concurrent_context_operations(self):
+        """Test concurrent context operations from multiple threads."""
+        results = []
+        errors = []
+
+        def thread_context_operation(thread_id):
+            try:
+                bst.environ.set(**{f'base_{thread_id}': 0})
+
+                for i in range(5):
+                    with bst.environ.context(**{f'base_{thread_id}': i}):
+                        value = bst.environ.get(f'base_{thread_id}')
+                        results.append((thread_id, value))
+
+                # Should be back to 0
+                final = bst.environ.get(f'base_{thread_id}')
+                self.assertEqual(final, 0)
+            except Exception as e:
+                errors.append(e)
+
+        threads = []
+        for i in range(3):
+            thread = threading.Thread(target=thread_context_operation, args=(i,))
+            threads.append(thread)
+            thread.start()
+
+        for thread in threads:
+            thread.join()
+
+        # Should have no errors
+        self.assertEqual(len(errors), 0)
+
+    def test_concurrent_pop_operations(self):
+        """Test concurrent pop operations from multiple threads."""
+        # Set up multiple keys
+        for i in range(20):
+            bst.environ.set(**{f'pop_thread_{i}': f'value_{i}'})
+
+        results = []
+        errors = []
+
+        def thread_pop_operation(start, end):
+            try:
+                for i in range(start, end):
+                    try:
+                        value = bst.environ.pop(f'pop_thread_{i}')
+                        results.append((i, value))
+                    except KeyError:
+                        # Key might already be popped by another thread
+                        pass
+            except Exception as e:
+                errors.append(e)
+
+        # Create threads that pop different ranges
+        threads = []
+        ranges = [(0, 5), (5, 10), (10, 15), (15, 20)]
+        for start, end in ranges:
+            thread = threading.Thread(target=thread_pop_operation, args=(start, end))
+            threads.append(thread)
+            thread.start()
+
+        for thread in threads:
+            thread.join()
+
+        # Should have no errors
+        self.assertEqual(len(errors), 0)
+
+        # # All keys should be popped (each exactly once)
+        # popped_indices = [r[0] for r in results]
+        # self.assertEqual(len(popped_indices), 20)
+        # self.assertEqual(len(set(popped_indices)), 20)  # All unique
+        #
+        # # All values should be gone
+        # for i in range(20):
+        #     result = brainstate.environ.get(f'pop_thread_{i}', default=None)
+        #     self.assertIsNone(result)
+
+
+class TestEdgeCases(unittest.TestCase):
+    """Test edge cases and boundary conditions."""
+
+    def setUp(self):
+        """Reset environment before each test."""
+        bst.environ.reset()
+        warnings.filterwarnings('ignore', category=UserWarning)
+
+    def tearDown(self):
+        """Clean up after each test."""
+        bst.environ.reset()
+        warnings.resetwarnings()
+
+    def test_empty_context(self):
+        """Test context with no parameters."""
+        original = bst.environ.all()
+
+        with bst.environ.context() as ctx:
+            # Should be unchanged
+            self.assertEqual(ctx, original)
+
+        self.assertEqual(bst.environ.all(), original)
+
+    def test_none_values(self):
+        """Test handling of None values."""
+        bst.environ.set(none_param=None)
+        self.assertIsNone(bst.environ.get('none_param'))
+
+        with bst.environ.context(none_param='not_none'):
+            self.assertEqual(bst.environ.get('none_param'), 'not_none')
+
+        self.assertIsNone(bst.environ.get('none_param'))
+
+    def test_complex_data_types(self):
+        """Test storing complex data types."""
+        # Lists
+        bst.environ.set(list_param=[1, 2, 3])
+        self.assertEqual(bst.environ.get('list_param'), [1, 2, 3])
+
+        # Dictionaries
+        bst.environ.set(dict_param={'a': 1, 'b': 2})
+        self.assertEqual(bst.environ.get('dict_param'), {'a': 1, 'b': 2})
+
+        # Tuples
+        bst.environ.set(tuple_param=(1, 2, 3))
+        self.assertEqual(bst.environ.get('tuple_param'), (1, 2, 3))
+
+        # Custom objects
+        class CustomObject:
+            def __init__(self, value):
+                self.value = value
+
+        obj = CustomObject(42)
+        bst.environ.set(obj_param=obj)
+        retrieved = bst.environ.get('obj_param')
+        self.assertIs(retrieved, obj)
+        self.assertEqual(retrieved.value, 42)
+
+    def test_special_string_values(self):
+        """Test special string values."""
+        special_strings = ['', ' ', '\n', '\t', 'None', 'True', 'False']
+
+        for s in special_strings:
+            bst.environ.set(string_param=s)
+            self.assertEqual(bst.environ.get('string_param'), s)
+
+    def test_numeric_edge_values(self):
+        """Test numeric edge values."""
+        import sys
+
+        edge_values = [
+            0, -0, 1, -1,
+            sys.maxsize, -sys.maxsize,
+            float('inf'), float('-inf'),
+            1e-100, 1e100,
+        ]
+
+        for value in edge_values:
+            bst.environ.set(numeric_param=value)
+            retrieved = bst.environ.get('numeric_param')
+            if value != value:  # NaN check
+                self.assertTrue(retrieved != retrieved)
+            else:
+                self.assertEqual(retrieved, value)
+
+    def test_context_all_interaction(self):
+        """Test interaction between context and all() function."""
+        bst.environ.set(global_param='global')
+
+        with bst.environ.context(context_param='context', global_param='override'):
+            all_values = bst.environ.all()
+
+            # Should include both
+            self.assertEqual(all_values['global_param'], 'override')
+            self.assertEqual(all_values['context_param'], 'context')
+
+            # Original global values should be in settings
+            self.assertIn('precision', all_values)
+
+    def test_deeply_nested_contexts(self):
+        """Test deeply nested contexts."""
+        depth = 20
+        bst.environ.set(depth=0)
+
+        def nested_context(level):
+            if level < depth:
+                with bst.environ.context(depth=level):
+                    self.assertEqual(bst.environ.get('depth'), level)
+                    nested_context(level + 1)
+                    self.assertEqual(bst.environ.get('depth'), level)
+
+        nested_context(1)
+        self.assertEqual(bst.environ.get('depth'), 0)
+
+    def test_set_precision_function(self):
+        """Test the dedicated set_precision function."""
+        # Valid precisions
+        for precision in [8, 16, 32, 64, 'bf16']:
+            bst.environ.set_precision(precision)
+            self.assertEqual(bst.environ.get('precision'), precision)
+
+        # Invalid precision
+        with self.assertRaises(ValueError):
+            bst.environ.set_precision(128)
+
+    def test_pop_edge_cases(self):
+        """Test edge cases for pop function."""
+        # Pop with None value
+        bst.environ.set(none_key=None)
+        popped = bst.environ.pop('none_key')
+        self.assertIsNone(popped)
+
+        # Pop with None as default
+        result = bst.environ.pop('missing_key', default=None)
+        self.assertIsNone(result)
+
+        # Pop complex data types
+        complex_obj = {'nested': {'data': [1, 2, 3]}}
+        bst.environ.set(complex_key=complex_obj)
+        popped = bst.environ.pop('complex_key')
+        self.assertEqual(popped, complex_obj)
+
+        # Verify object identity preservation
+        obj = object()
+        bst.environ.set(obj_key=obj)
+        popped = bst.environ.pop('obj_key')
+        self.assertIs(popped, obj)
+
+    def test_pop_all_interaction(self):
+        """Test interaction between pop and all() function."""
+        # Set multiple values
+        bst.environ.set(a=1, b=2, c=3, d=4)
+        initial_all = bst.environ.all()
+
+        # Pop some values
+        bst.environ.pop('b')
+        bst.environ.pop('d')
+
+        # Check all() reflects the changes
+        after_pop = bst.environ.all()
+        self.assertIn('a', after_pop)
+        self.assertIn('c', after_pop)
+        self.assertNotIn('b', after_pop)
+        self.assertNotIn('d', after_pop)
+
+    def test_pop_callback_not_triggered(self):
+        """Test that callbacks are not triggered on pop."""
+        callback_calls = []
+
+        def callback(value):
+            callback_calls.append(value)
+
+        # Register callback
+        bst.environ.register_default_behavior('callback_test', callback)
+
+        # Set triggers callback
+        bst.environ.set(callback_test='value')
+        self.assertEqual(len(callback_calls), 1)
+
+        # Pop should NOT trigger callback
+        popped = bst.environ.pop('callback_test')
+        self.assertEqual(len(callback_calls), 1)  # Still just 1
+        self.assertEqual(popped, 'value')
+
+        # Unregister callback
+        bst.environ.unregister_default_behavior('callback_test')
+
+
+class TestIntegration(unittest.TestCase):
+    """Integration tests with actual BrainState functionality."""
+
+    def setUp(self):
+        """Reset environment before each test."""
+        bst.environ.reset()
+        warnings.filterwarnings('ignore', category=UserWarning)
+
+    def tearDown(self):
+        """Clean up after each test."""
+        bst.environ.reset()
+        warnings.resetwarnings()
+
+    def test_precision_affects_random_arrays(self):
+        """Test that precision setting affects random array generation."""
+        # Test different precisions
+        test_cases = [
+            (32, jnp.float32),
+            (64, jnp.float64),
+            (16, jnp.float16),
+            ('bf16', jnp.bfloat16),
+        ]
+
+        for precision, expected_dtype in test_cases:
+            with bst.environ.context(precision=precision):
+                arr = bst.random.randn(10)
+                self.assertEqual(arr.dtype, expected_dtype)
+
+    def test_mode_usage(self):
+        """Test mode usage in computations."""
+        # Create different modes
+        training = bst.mixin.Training()
+        batching = bst.mixin.Batching(batch_size=32)
+
+        # Test training mode
+        bst.environ.set(mode=training)
+        mode = bst.environ.get('mode')
+        self.assertTrue(mode.has(bst.mixin.Training))
+
+        # Test batching mode
+        with bst.environ.context(mode=batching):
+            mode = bst.environ.get('mode')
+            self.assertTrue(mode.has(bst.mixin.Batching))
+            self.assertEqual(mode.batch_size, 32)
+
+    def test_dt_in_numerical_integration(self):
+        """Test dt usage in numerical contexts."""
+        # Set different dt values
+        dt_values = [0.01, 0.001, 0.1]
+
+        for dt in dt_values:
+            bst.environ.set(dt=dt)
+            retrieved_dt = bst.environ.get_dt()
+            self.assertEqual(retrieved_dt, dt)
+
+            # Simulate using dt in computation
+            time_steps = int(1.0 / dt)
+            self.assertGreater(time_steps, 0)
+
+    def test_combined_settings(self):
+        """Test combining multiple settings."""
+        # Set multiple parameters
+        bst.environ.set(
+            precision=64,
+            dt=0.01,
+            mode=bst.mixin.Training(),
+            custom_param='test',
+            debug=True
+        )
+
+        # Verify all are set
+        self.assertEqual(bst.environ.get_precision(), 64)
+        self.assertEqual(bst.environ.get_dt(), 0.01)
+        self.assertTrue(bst.environ.get('mode').has(bst.mixin.Training))
+        self.assertEqual(bst.environ.get('custom_param'), 'test')
+        self.assertTrue(bst.environ.get('debug'))
+
+        # Test in nested contexts
+        with bst.environ.context(precision=32, debug=False):
+            self.assertEqual(bst.environ.get_precision(), 32)
+            self.assertFalse(bst.environ.get('debug'))
+            # Others unchanged
+            self.assertEqual(bst.environ.get_dt(), 0.01)
+            self.assertEqual(bst.environ.get('custom_param'), 'test')
+
+
+if __name__ == '__main__':
+    unittest.main()