Spaces:
Sleeping
Sleeping
| """ | |
| Unit tests for PyTorch integration in CloudResourceEnv. | |
| This module tests the minimal PyTorch tensor-based state representation | |
| added to the environment while ensuring NumPy compatibility is maintained. | |
| """ | |
| import pytest | |
| import numpy as np | |
| import torch | |
| from env.environment import CloudResourceEnv | |
| from env.config import EnvConfig | |
| def test_get_state_tensor_returns_torch_tensor(): | |
| """Test that get_state_tensor returns a PyTorch tensor.""" | |
| env = CloudResourceEnv() | |
| env.reset() | |
| state_tensor = env.get_state_tensor() | |
| assert isinstance(state_tensor, torch.Tensor), "get_state_tensor should return a torch.Tensor" | |
| def test_get_state_tensor_matches_numpy_observation(): | |
| """Test that PyTorch tensor state matches NumPy observation.""" | |
| env = CloudResourceEnv() | |
| observation = env.reset() | |
| state_tensor = env.get_state_tensor() | |
| # Convert tensor to numpy for comparison | |
| tensor_as_numpy = state_tensor.numpy() | |
| # Should match the observation exactly | |
| np.testing.assert_array_equal( | |
| tensor_as_numpy, | |
| observation, | |
| err_msg="PyTorch tensor state should match NumPy observation" | |
| ) | |
| def test_get_state_tensor_after_step(): | |
| """Test that get_state_tensor works correctly after environment steps.""" | |
| env = CloudResourceEnv() | |
| env.reset() | |
| # Take a step | |
| observation, reward, done, info = env.step(CloudResourceEnv.ACTION_INCREASE) | |
| # Get tensor representation | |
| state_tensor = env.get_state_tensor() | |
| # Verify it matches the observation | |
| tensor_as_numpy = state_tensor.numpy() | |
| np.testing.assert_array_equal( | |
| tensor_as_numpy, | |
| observation, | |
| err_msg="Tensor state should match observation after step" | |
| ) | |
| def test_get_state_tensor_shape(): | |
| """Test that state tensor has correct shape.""" | |
| env = CloudResourceEnv() | |
| env.reset() | |
| state_tensor = env.get_state_tensor() | |
| assert state_tensor.shape == (4,), f"State tensor should have shape (4,), got {state_tensor.shape}" | |
| def test_get_state_tensor_dtype(): | |
| """Test that state tensor has correct dtype (float32).""" | |
| env = CloudResourceEnv() | |
| env.reset() | |
| state_tensor = env.get_state_tensor() | |
| assert state_tensor.dtype == torch.float32, f"State tensor should be float32, got {state_tensor.dtype}" | |
| def test_numpy_compatibility_maintained(): | |
| """Test that existing NumPy-based functionality still works.""" | |
| env = CloudResourceEnv() | |
| # Reset should still return NumPy array | |
| observation = env.reset() | |
| assert isinstance(observation, np.ndarray), "reset() should still return NumPy array" | |
| # Step should still return NumPy array | |
| observation, reward, done, info = env.step(CloudResourceEnv.ACTION_MAINTAIN) | |
| assert isinstance(observation, np.ndarray), "step() should still return NumPy array" | |
| def test_get_state_tensor_multiple_calls(): | |
| """Test that multiple calls to get_state_tensor return consistent results.""" | |
| env = CloudResourceEnv() | |
| env.reset() | |
| tensor1 = env.get_state_tensor() | |
| tensor2 = env.get_state_tensor() | |
| # Should be equal since state hasn't changed | |
| assert torch.equal(tensor1, tensor2), "Multiple calls should return equal tensors" | |
| def test_get_state_tensor_updates_with_state(): | |
| """Test that state tensor updates when environment state changes.""" | |
| env = CloudResourceEnv() | |
| env.reset() | |
| tensor_before = env.get_state_tensor().clone() | |
| # Take multiple steps to ensure state changes | |
| for _ in range(5): | |
| observation, reward, done, info = env.step(CloudResourceEnv.ACTION_INCREASE) | |
| if done: | |
| break | |
| tensor_after = env.get_state_tensor() | |
| # Tensors should be different (state has changed) | |
| assert not torch.equal(tensor_before, tensor_after), "State tensor should update with environment state" | |