engine/test_indicators.py

import os
import json
import unittest
from pathlib import Path
from unittest.mock import patch, MagicMock

# Import the modules to test
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from memory_management import MemoryManager
import asset_generator

class TestIndicators(unittest.TestCase):
    def test_memory_manager_status(self):
        MemoryManager.set_status("TESTING")
        self.assertEqual(MemoryManager.status, "TESTING")
        MemoryManager.set_status("IDLE")
        self.assertEqual(MemoryManager.status, "IDLE")

    @patch('subprocess.check_output')
    def test_get_gpu_usage(self, mock_subprocess):
        # Mock nvidia-smi output
        mock_subprocess.return_value = "42\n"
        usage = asset_generator.get_gpu_usage()
        self.assertEqual(usage, 42)

        # Mock failure
        mock_subprocess.side_effect = Exception("nvidia-smi not found")
        usage = asset_generator.get_gpu_usage()
        self.assertEqual(usage, 0)

    def test_current_task_json_format(self):
        # Test if the dictionary format matches what's expected by the dashboard
        test_data = {
            "story": "test_story",
            "story_id": "test_id",
            "progress": 50,
            "img": "local",
            "voc": "api",
            "mem_status": "IDLE"
        }
        # This is just to ensure no typos in keys if we were to refactor
        required_keys = ["story", "story_id", "progress", "img", "voc", "mem_status"]
        for key in required_keys:
            self.assertIn(key, test_data)

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