# Copyright (c) Huawei Technologies Co., Ltd. 2026-2026. All rights reserved.
# MindIE is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         `http://license.coscl.org.cn/MulanPSL2`
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

import unittest
import threading
import time
from unittest.mock import Mock, patch

from olc.collector.collector import AbsCollector, Meter, Measurable, CollectorTask

class MockMeasurable(Measurable):
    @classmethod
    def get_properties(cls) -> list:
        return ["value", "count"]


class MockMeter(Meter):
    def __init__(self, value=0):
        self.value = value

    def measure(self):
        return MockMeasurement(self.value)


class MockMeasurement:
    def __init__(self, value):
        self.value = value
        self.count = 1


class MockCollector(AbsCollector):
    def __init__(self, meter=None, measurable_type=MockMeasurable, **kwargs):
        super().__init__(meter=meter, measurable_type=measurable_type, **kwargs)
        self._storage = []

    def get_name(self) -> str:
        return "test-collector"

    def get_storage(self) -> list:
        return self._storage


class TestCollectorTask(unittest.TestCase):
    def test_run(self):
        collector = Mock(spec=AbsCollector)
        task = CollectorTask(collector)
        task.run()
        collector.do_measure.assert_called_once()

    def test_get_name(self):
        collector = Mock(spec=AbsCollector)
        collector.get_name.return_value = "test-collector"
        task = CollectorTask(collector)
        self.assertEqual(task.get_name(), "test-collector-task")


class TestAbsCollector(unittest.TestCase):
    def setUp(self):
        # 模拟配置
        with patch('olc.collector.collector.OlcConfigManager') as mock_config:
            mock_instance = Mock()
            mock_sdk_config = Mock()
            mock_sdk_config.collector_interval = 1
            mock_sdk_config.collector_capacity = 100
            mock_sdk_config.collector_sample = 10
            mock_instance.get_sdk_config.return_value = mock_sdk_config
            mock_config.get_instance.return_value = mock_instance

            self.collector = MockCollector(meter=MockMeter(5), should_start=False)

    def test_start(self):
        with patch('olc.collector.collector.OlcScheduleManager') as mock_schedule:
            mock_instance = Mock()
            mock_schedule.get_instance.return_value = mock_instance

            self.collector.start()
            mock_instance.create_task.assert_called_once()
            mock_instance.start_task.assert_called_once()
            self.assertTrue(self.collector._is_running)

    def test_stop(self):
        with patch('olc.collector.collector.OlcScheduleManager') as mock_schedule:
            mock_instance = Mock()
            mock_schedule.get_instance.return_value = mock_instance

            self.collector.start()
            self.collector.stop()
            mock_instance.stop_task.assert_called_once()

    def test_do_measure(self):
        self.collector.do_measure()
        storage = self.collector.get_storage()
        self.assertEqual(len(storage), 1)
        self.assertEqual(storage[0].value, 5)

    def test_do_measure_with_none_result(self):
        meter = Mock(spec=Meter)
        meter.measure.return_value = None
        self.collector._meter = meter

        self.collector.do_measure()
        storage = self.collector.get_storage()
        self.assertEqual(len(storage), 0)

    def test_do_measure_with_exception(self):
        meter = Mock(spec=Meter)
        meter.measure.side_effect = Exception("Test exception")
        self.collector._meter = meter

        # 确保不会抛出异常
        self.collector.do_measure()

    def test_get_properties(self):
        properties = self.collector.get_properties()
        self.assertEqual(properties, ["value", "count"])

    def test_get_average_property(self):
        # 添加测试数据
        for i in range(5):
            self.collector._storage.append(MockMeasurement(i))

        average = self.collector.get_average_property("value")
        self.assertAlmostEqual(average, 2.0)

    def test_get_average_property_no_data(self):
        average = self.collector.get_average_property("value")
        self.assertIsNone(average)

    def test_get_average_property_invalid_property(self):
        self.collector._storage.append(MockMeasurement(5))
        average = self.collector.get_average_property("invalid")
        self.assertIsNone(average)

    def test_get_average_property_non_float(self):
        # 创建一个非float属性的测量值
        class NonFloatMeasurement:
            def __init__(self):
                self.value = "string"  # 非float类型

        self.collector._storage.append(NonFloatMeasurement())
        average = self.collector.get_average_property("value")
        self.assertIsNone(average)

    def test_storage_truncation(self):
        # 设置样本数为3
        self.collector._capacity = 3

        # 添加5个测量值
        for i in range(5):
            self.collector._storage.append(MockMeasurement(i))

        # 执行测量，应该截断到最近的3个，然后添加新的测量值
        self.collector.do_measure()
        storage = self.collector.get_storage()
        self.assertEqual(len(storage), 4)  # 3个截断后的值 + 1个新测量值


class TestCollectorConcurrent(unittest.TestCase):
    def setUp(self):
        with patch('olc.collector.collector.OlcConfigManager') as mock_config:
            mock_instance = Mock()
            mock_sdk_config = Mock()
            mock_sdk_config.collector_interval = 1
            mock_sdk_config.collector_capacity = 100
            mock_instance.get_sdk_config.return_value = mock_sdk_config
            mock_config.get_instance.return_value = mock_instance

            self.collector = MockCollector(meter=MockMeter(1), should_start=False)

    def test_concurrent_measure(self):
        def measure_task():
            for _ in range(10):
                self.collector.do_measure()
                time.sleep(0.01)

        # 创建多个线程同时执行测量
        threads = []
        for _ in range(5):
            t = threading.Thread(target=measure_task)
            threads.append(t)
            t.start()

        # 等待所有线程完成
        for t in threads:
            t.join()

        # 验证存储中有数据
        storage = self.collector.get_storage()
        self.assertTrue(len(storage) > 0)


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