water-management-system/tests/unit/test_data_engine.py

"""
数据引擎单元测试
覆盖CRUD操作、批量导入功能
"""
import unittest
from unittest.mock import Mock, patch, MagicMock
import json
from src.data.engine import DataEngine
from src.data.models import DataModel, BatchImportModel


class TestDataEngine(unittest.TestCase):
    """数据引擎测试"""
    
    def setUp(self):
        self.engine = DataEngine("postgresql://localhost:5432/water_db")
    
    def test_create_operation(self):
        """测试数据创建"""
        test_data = {
            "device_id": "test_device_001",
            "timestamp": "2026-06-16T12:00:00Z",
            "temperature": 25.5,
            "humidity": 60.2
        }
        
        with patch('src.data.engine.session') as mock_session:
            mock_query = Mock()
            mock_query.return_value = Mock()
            mock_query.return_value.execute.return_value = Mock()
            mock_session.query.return_value = mock_query
            
            result = self.engine.create("sensor_data", test_data)
            self.assertTrue(result["success"])
            mock_session.add.assert_called_once()
    
    def test_read_operation(self):
        """测试数据读取"""
        device_id = "test_device_001"
        
        with patch('src.data.engine.session') as mock_session:
            mock_query = Mock()
            mock_session.query.return_value.filter.return_value.all.return_value = [
                {"device_id": device_id, "temperature": 25.5}
            ]
            
            results = self.engine.read("sensor_data", {"device_id": device_id})
            self.assertEqual(len(results), 1)
            self.assertEqual(results[0]["device_id"], device_id)
    
    def test_update_operation(self):
        """测试数据更新"""
        update_data = {"temperature": 26.0}
        condition = {"device_id": "test_device_001"}
        
        with patch('src.data.engine.session') as mock_session:
            mock_query = Mock()
            mock_session.query.return_value.filter.return_value.update.return_value = 1
            
            result = self.engine.update("sensor_data", update_data, condition)
            self.assertEqual(result["updated_count"], 1)
    
    def test_delete_operation(self):
        """测试数据删除"""
        condition = {"device_id": "test_device_001", "timestamp": "2026-06-16T12:00:00Z"}
        
        with patch('src.data.engine.session') as mock_session:
            mock_query = Mock()
            mock_session.query.return_value.filter.return_value.delete.return_value = 1
            
            result = self.engine.delete("sensor_data", condition)
            self.assertEqual(result["deleted_count"], 1)
    
    def test_batch_import(self):
        """测试批量导入"""
        batch_data = [
            {"device_id": "dev001", "temperature": 25.0, "timestamp": "2026-06-16T12:00:00Z"},
            {"device_id": "dev002", "temperature": 26.1, "timestamp": "2026-06-16T12:01:00Z"},
            {"device_id": "dev003", "temperature": 24.8, "timestamp": "2026-06-16T12:02:00Z"}
        ]
        
        with patch('src.data.engine.session') as mock_session:
            mock_batch_model = Mock()
            mock_session.bulk_save_objects.return_value = Mock()
            
            result = self.engine.batch_import("sensor_data", batch_data)
            self.assertEqual(result["imported_count"], 3)
            mock_session.bulk_save_objects.assert_called_once()
    
    def test_batch_import_validation(self):
        """测试批量导入数据验证"""
        batch_data = [
            {"device_id": "dev001", "temperature": "invalid_temp"},  # 无效温度
            {"device_id": "dev002", "temperature": 26.1}
        ]
        
        with patch('src.data.engine.DataModel.validate') as mock_validate:
            mock_validate.side_effect = [
                {"valid": False, "errors": ["Invalid temperature format"]},
                {"valid": True, "errors": []}
            ]
            
            result = self.engine.batch_import("sensor_data", batch_data)
            self.assertEqual(result["imported_count"], 1)
            self.assertEqual(result["validation_errors"], 1)


class TestDataModel(unittest.TestCase):
    """数据模型测试"""
    
    def setUp(self):
        self.model = DataModel()
    
    def test_field_validation(self):
        """测试字段验证"""
        valid_data = {
            "device_id": "test_device",
            "timestamp": "2026-06-16T12:00:00Z",
            "temperature": 25.5,
            "humidity": 60.2
        }
        
        result = self.model.validate(valid_data)
        self.assertTrue(result["valid"])
    
    def test_invalid_data_validation(self):
        """测试无效数据验证"""
        invalid_data = {
            "device_id": "",  # 空设备ID
            "temperature": "not_a_number",  # 非数字温度
            "timestamp": "invalid_timestamp"  # 无效时间戳
        }
        
        result = self.model.validate(invalid_data)
        self.assertFalse(result["valid"])
        self.assertGreater(len(result["errors"]), 0)


class TestBatchImportModel(unittest.TestCase):
    """批量导入模型测试"""
    
    def setUp(self):
        self.batch_model = BatchImportModel()
    
    def test_batch_size_limit(self):
        """测试批量大小限制"""
        # 创建超过限制的数据
        large_batch = [{"device_id": f"dev{i}", "value": i} for i in range(1001)]
        
        result = self.batch_model.validate_batch(large_batch)
        self.assertFalse(result["valid"])
        self.assertIn("Batch size exceeds limit", result["errors"])
    
    def test_batch_structure_validation(self):
        """测试批量结构验证"""
        valid_batch = [
            {"device_id": "dev001", "value": 100},
            {"device_id": "dev002", "value": 200}
        ]
        
        result = self.batch_model.validate_batch(valid_batch)
        self.assertTrue(result["valid"])


if __name__ == '__main__':
    unittest.main()