import unittest
from unittest.mock import patch
from aiva.cpu_frequency_governor import CPUFrequencyGovernor, HYSTERESIS_WINDOW

class TestCPUFrequencyGovernor(unittest.TestCase):

    def setUp(self):
        self.governor = CPUFrequencyGovernor()
        self.governor.available_frequencies = [1000, 2000, 3000]
        self.governor.current_frequency = 2000
        self.governor.high_utilization_threshold = 70
        self.governor.low_utilization_threshold = 30

    def test_scale_up_hysteresis(self):
        # Simulate sustained high utilization above the threshold
        with patch('aiva.cpu_frequency_governor.get_cpu_utilization', return_value=80) as mock_utilization:
            for _ in range(HYSTERESIS_WINDOW + 1):
                self.governor.adjust_frequency()
        
        self.assertEqual(self.governor.current_frequency, 3000, "Frequency should scale up after sustained high utilization")

    def test_scale_down_hysteresis(self):
        # First, scale up to max frequency
        self.governor.current_frequency = 3000
        
        # Simulate sustained low utilization below the threshold
        with patch('aiva.cpu_frequency_governor.get_cpu_utilization', return_value=20) as mock_utilization:
            for _ in range(HYSTERESIS_WINDOW + 1):
                self.governor.adjust_frequency()

        self.assertEqual(self.governor.current_frequency, 2000, "Frequency should scale down after sustained low utilization")

    def test_no_scale_up_before_hysteresis(self):
        # Simulate utilization slightly above the threshold
        with patch('aiva.cpu_frequency_governor.get_cpu_utilization', return_value=75) as mock_utilization:
            for _ in range(HYSTERESIS_WINDOW):
                self.governor.adjust_frequency()

        self.assertEqual(self.governor.current_frequency, 2000, "Frequency should not scale up before hysteresis window")

    def test_no_scale_down_before_hysteresis(self):
        # First, scale up to max frequency
        self.governor.current_frequency = 3000

        # Simulate utilization slightly below the threshold
        with patch('aiva.cpu_frequency_governor.get_cpu_utilization', return_value=25) as mock_utilization:
            for _ in range(HYSTERESIS_WINDOW):
                self.governor.adjust_frequency()

        self.assertEqual(self.governor.current_frequency, 3000, "Frequency should not scale down before hysteresis window")

    def test_scale_up_and_down_alternating(self):
        # Simulate alternating high and low utilization
        with patch('aiva.cpu_frequency_governor.get_cpu_utilization') as mock_utilization:
            # High Utilization
            mock_utilization.return_value = 80
            for _ in range(HYSTERESIS_WINDOW + 1):
                self.governor.adjust_frequency()

            self.assertEqual(self.governor.current_frequency, 3000, "Frequency should scale up")

            # Low Utilization
            mock_utilization.return_value = 20
            for _ in range(HYSTERESIS_WINDOW + 1):
                self.governor.adjust_frequency()

            self.assertEqual(self.governor.current_frequency, 2000, "Frequency should scale down")

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