#!/usr/bin/python3
# -*- coding: utf-8 -*-
# SPDX-License-Identifier: GPL-2.0-only
"""
Comprehensive test suite for tuna.cpuset module

Tests cover:
- Core cpuset functionality (rteval-compatible)
- Tuna-specific enhancements (discovery, cleanup, persistence)
- Edge cases and error handling
- Nested cpuset hierarchies

Requires root permissions to run tests that create/destroy cgroups.
"""

import unittest
import os
import sys
import time
import tempfile
from pathlib import Path

# Add tuna to path if running from tests directory
if Path(__file__).parent.name == 'tests':
    sys.path.insert(0, str(Path(__file__).parent.parent))

from tuna import cpuset


class TestCpusetsInit(unittest.TestCase):
    """Test CpusetsInit class and cgroup v2 detection"""

    def test_initialization(self):
        """Test that CpusetsInit initializes correctly"""
        ci = cpuset.CpusetsInit()
        self.assertIsInstance(ci, cpuset.CpusetsInit)
        self.assertIsInstance(ci.supported, bool)
        self.assertIsInstance(ci.numa_nodes, int)

    def test_cpuset_path(self):
        """Test that cpuset_path is set correctly for root cgroup"""
        ci = cpuset.CpusetsInit()
        if ci.supported:
            self.assertEqual(ci.cpuset_path, '/sys/fs/cgroup')

    def test_cgroup_v2_detection(self):
        """Test cgroup v2 detection logic"""
        supported = cpuset.CpusetsInit._cpuset_supported()
        self.assertIsInstance(supported, bool)

        # If /sys/fs/cgroup exists and has cgroup.controllers, should be supported
        if os.path.exists('/sys/fs/cgroup/cgroup.controllers'):
            # Read controllers to verify cpuset is available
            with open('/sys/fs/cgroup/cgroup.controllers') as f:
                controllers = f.read().strip().split()
                if 'cpuset' in controllers:
                    self.assertTrue(supported)

    def test_numa_nodes(self):
        """Test NUMA node detection"""
        ci = cpuset.CpusetsInit()
        if ci.supported:
            # Should have at least 1 NUMA node
            self.assertGreaterEqual(ci.numa_nodes, 1)

            # Verify against actual /sys/devices/system/node
            from glob import glob
            actual_nodes = len(glob('/sys/devices/system/node/node*'))
            self.assertEqual(ci.numa_nodes, actual_nodes)

    def test_auto_destroy_is_false(self):
        """Test that root cgroup never has auto_destroy enabled"""
        ci = cpuset.CpusetsInit()
        if ci.supported:
            self.assertFalse(ci.auto_destroy)


@unittest.skipUnless(os.geteuid() == 0, "Requires root permissions")
@unittest.skipUnless(cpuset.CpusetsInit().supported, "Requires cgroup v2 support")
class TestCpusetCreation(unittest.TestCase):
    """Test Cpuset creation and destruction"""

    def setUp(self):
        """Set up test fixtures"""
        self.test_cpusets = []
        self.ci = cpuset.CpusetsInit()

    def tearDown(self):
        """Clean up any cpusets created during tests"""
        for cs in self.test_cpusets:
            try:
                if hasattr(cs, 'cpuset_path') and os.path.exists(cs.cpuset_path):
                    # Move any tasks back to root
                    tm = cpuset.TaskMigrate(cs, self.ci)
                    tm.migrate()
                    cs.destroy()
            except:
                pass

    def test_create_cpuset(self):
        """Test basic cpuset creation"""
        cs = cpuset.Cpuset('test_create')
        self.test_cpusets.append(cs)

        self.assertEqual(cs.cpuset_name, 'test_create')
        self.assertEqual(cs.cpuset_path, '/sys/fs/cgroup/test_create')
        self.assertTrue(os.path.exists(cs.cpuset_path))

    def test_create_with_auto_destroy_false(self):
        """Test cpuset creation with auto_destroy=False (default)"""
        cs = cpuset.Cpuset('test_no_auto')
        self.test_cpusets.append(cs)

        self.assertFalse(cs.auto_destroy)

    def test_create_with_auto_destroy_true(self):
        """Test cpuset creation with auto_destroy=True"""
        cs = cpuset.Cpuset('test_auto', auto_destroy=True)
        self.test_cpusets.append(cs)

        self.assertTrue(cs.auto_destroy)

    def test_destroy_cpuset(self):
        """Test manual cpuset destruction"""
        cs = cpuset.Cpuset('test_destroy')
        path = cs.cpuset_path

        self.assertTrue(os.path.exists(path))
        cs.destroy()
        self.assertFalse(os.path.exists(path))

    def test_create_existing_cpuset(self):
        """Test creating a cpuset that already exists"""
        cs1 = cpuset.Cpuset('test_existing')
        self.test_cpusets.append(cs1)
        path1 = cs1.cpuset_path

        # Create again with same name
        cs2 = cpuset.Cpuset('test_existing')

        # Should reference the same path
        self.assertEqual(cs1.cpuset_path, cs2.cpuset_path)
        self.assertTrue(os.path.exists(path1))

    def test_cpuset_name_none_raises_error(self):
        """Test that creating cpuset with None name raises ValueError"""
        with self.assertRaises(ValueError):
            cpuset.Cpuset(None)


@unittest.skipUnless(os.geteuid() == 0, "Requires root permissions")
@unittest.skipUnless(cpuset.CpusetsInit().supported, "Requires cgroup v2 support")
class TestCpusetConfiguration(unittest.TestCase):
    """Test Cpuset configuration (CPUs, memory, etc.)"""

    def setUp(self):
        """Set up test fixtures"""
        self.test_cpusets = []
        self.ci = cpuset.CpusetsInit()

    def tearDown(self):
        """Clean up any cpusets created during tests"""
        for cs in self.test_cpusets:
            try:
                if hasattr(cs, 'cpuset_path') and os.path.exists(cs.cpuset_path):
                    tm = cpuset.TaskMigrate(cs, self.ci)
                    tm.migrate()
                    cs.destroy()
            except:
                pass

    def test_assign_cpus(self):
        """Test CPU assignment to cpuset"""
        cs = cpuset.Cpuset('test_cpus')
        self.test_cpusets.append(cs)

        cs.write_memnode('0')
        cs.assign_cpus('0-1')

        # Read back the assignment
        with open(os.path.join(cs.cpuset_path, 'cpuset.cpus')) as f:
            assigned = f.read().strip()

        # Should be '0-1' or '0,1' depending on kernel
        self.assertIn('0', assigned)
        self.assertIn('1', assigned)

    def test_write_memnode(self):
        """Test memory node assignment"""
        cs = cpuset.Cpuset('test_memnode')
        self.test_cpusets.append(cs)

        cs.write_memnode('0')

        # Read back the assignment
        with open(os.path.join(cs.cpuset_path, 'cpuset.mems')) as f:
            assigned = f.read().strip()

        self.assertEqual(assigned, '0')

    def test_cpu_exclusive(self):
        """Test CPU exclusive (partition) setting"""
        # Find a CPU that isn't already exclusively allocated to avoid conflicts
        # Read all existing cpusets and their exclusive CPUs
        used_cpus = set()
        all_cpusets = cpuset.list_cpusets(recursive=True)

        for cs_name in all_cpusets:
            cs_path = os.path.join('/sys/fs/cgroup', cs_name)
            partition_file = os.path.join(cs_path, 'cpuset.cpus.partition')

            # Check if this cpuset has an isolated partition
            if os.path.exists(partition_file):
                with open(partition_file) as f:
                    partition = f.read().strip()
                if partition.startswith('isolated'):
                    # This cpuset has exclusively allocated CPUs
                    cpus_file = os.path.join(cs_path, 'cpuset.cpus')
                    if os.path.exists(cpus_file):
                        with open(cpus_file) as f:
                            cpus_str = f.read().strip()
                        if cpus_str:
                            # Parse CPU list and add to used set
                            from tuna import tuna
                            try:
                                cpu_list = tuna.expand_cpulist(cpus_str)
                                used_cpus.update(cpu_list)
                            except:
                                pass

        # Get actual number of CPUs on this system
        # Read from root cgroup's effective CPUs
        root_cpus_file = '/sys/fs/cgroup/cpuset.cpus.effective'
        try:
            with open(root_cpus_file) as f:
                root_cpus_str = f.read().strip()
            from tuna import tuna
            available_cpus = tuna.expand_cpulist(root_cpus_str)
        except:
            # Fallback: use multiprocessing to get CPU count
            import multiprocessing
            available_cpus = list(range(multiprocessing.cpu_count()))

        # Find first available CPU that isn't exclusively allocated
        test_cpu = None
        for cpu_num in available_cpus:
            if cpu_num not in used_cpus:
                test_cpu = str(cpu_num)
                break

        if test_cpu is None:
            self.skipTest("No available CPUs found for exclusive partition test (all CPUs are exclusively allocated)")

        cs = cpuset.Cpuset('test_exclusive')
        self.test_cpusets.append(cs)

        cs.write_memnode('0')
        cs.assign_cpus(test_cpu)

        # Try to set exclusive (may not be supported on all kernels)
        cs.set_cpu_exclusive()

        partition_file = os.path.join(cs.cpuset_path, 'cpuset.cpus.partition')
        if os.path.exists(partition_file):
            with open(partition_file) as f:
                partition = f.read().strip()
            # Should be 'isolated', 'root', or 'member'
            self.assertIn(partition, ['isolated', 'root', 'member'],
                         f"Unexpected partition state: {partition}")

    def test_get_cpus(self):
        """Test getting current CPU assignment"""
        cs = cpuset.Cpuset('test_get_cpus')
        self.test_cpusets.append(cs)

        cs.write_memnode('0')
        cs.assign_cpus('0-1')

        # Get CPUs using new method
        cpus = cs.get_cpus()

        # Should return a non-empty string with '0' and '1'
        self.assertIsInstance(cpus, str)
        self.assertTrue(cpus)
        self.assertIn('0', cpus)
        self.assertIn('1', cpus)

    def test_get_memnode(self):
        """Test getting current memory node assignment"""
        cs = cpuset.Cpuset('test_get_memnode')
        self.test_cpusets.append(cs)

        cs.write_memnode('0')

        # Get memory node using new method
        mems = cs.get_memnode()

        self.assertEqual(mems, '0')

    def test_get_partition_type(self):
        """Test getting current partition type"""
        cs = cpuset.Cpuset('test_get_partition')
        self.test_cpusets.append(cs)

        cs.write_memnode('0')
        cs.assign_cpus('0')

        # Get partition type
        partition = cs.get_partition_type()

        # Should be 'member', 'isolated', 'root', or None
        self.assertIn(partition, ['member', 'isolated', 'root', None])

    def test_modify_cpus_add(self):
        """Test adding CPUs to a cpuset"""
        cs = cpuset.Cpuset('test_modify_add')
        self.test_cpusets.append(cs)

        cs.write_memnode('0')
        cs.assign_cpus('0-1')

        # Add CPUs 2-3
        cs.modify_cpus(add_cpus=[2, 3])

        # Check result
        from tuna import tuna
        cpus_str = cs.get_cpus()
        cpu_list = tuna.expand_cpulist(cpus_str)

        self.assertIn(0, cpu_list)
        self.assertIn(1, cpu_list)
        self.assertIn(2, cpu_list)
        self.assertIn(3, cpu_list)

    def test_modify_cpus_remove(self):
        """Test removing CPUs from a cpuset"""
        cs = cpuset.Cpuset('test_modify_remove')
        self.test_cpusets.append(cs)

        cs.write_memnode('0')
        cs.assign_cpus('0-3')

        # Remove CPUs 2-3
        cs.modify_cpus(remove_cpus=[2, 3])

        # Check result
        from tuna import tuna
        cpus_str = cs.get_cpus()
        cpu_list = tuna.expand_cpulist(cpus_str)

        self.assertIn(0, cpu_list)
        self.assertIn(1, cpu_list)
        self.assertNotIn(2, cpu_list)
        self.assertNotIn(3, cpu_list)

    def test_modify_cpus_add_and_remove(self):
        """Test adding and removing CPUs in a single operation"""
        cs = cpuset.Cpuset('test_modify_both')
        self.test_cpusets.append(cs)

        cs.write_memnode('0')
        cs.assign_cpus('0-3')

        # Remove 0-1, add 4-5
        cs.modify_cpus(add_cpus=[4, 5], remove_cpus=[0, 1])

        # Check result
        from tuna import tuna
        cpus_str = cs.get_cpus()
        cpu_list = tuna.expand_cpulist(cpus_str)

        self.assertNotIn(0, cpu_list)
        self.assertNotIn(1, cpu_list)
        self.assertIn(2, cpu_list)
        self.assertIn(3, cpu_list)
        self.assertIn(4, cpu_list)
        self.assertIn(5, cpu_list)

    def test_modify_cpus_empty_result(self):
        """Test that modify_cpus handles removing all CPUs

        Note: The kernel may reject removing all CPUs if there are tasks
        in the cpuset (processes need CPUs to run on), so we just verify
        the method doesn't crash rather than asserting a specific result.
        """
        cs = cpuset.Cpuset('test_modify_empty')
        self.test_cpusets.append(cs)

        cs.write_memnode('0')
        cs.assign_cpus('0-1')

        # Try to remove all CPUs
        # This may or may not succeed depending on whether tasks are present
        cs.modify_cpus(remove_cpus=[0, 1])

        # Verify method completed without crashing
        # The kernel may keep the old CPU assignment if tasks are present
        cpus_str = cs.get_cpus()
        self.assertIsInstance(cpus_str, str)


@unittest.skipUnless(os.geteuid() == 0, "Requires root permissions")
@unittest.skipUnless(cpuset.CpusetsInit().supported, "Requires cgroup v2 support")
class TestContextManager(unittest.TestCase):
    """Test context manager behavior with auto_destroy"""

    def setUp(self):
        """Set up test fixtures"""
        self.ci = cpuset.CpusetsInit()

    def test_context_manager_auto_destroy_true(self):
        """Test that cpuset is destroyed when auto_destroy=True"""
        path = None

        with cpuset.Cpuset('test_context_auto', auto_destroy=True) as cs:
            path = cs.cpuset_path
            cs.write_memnode('0')
            cs.assign_cpus('0')
            self.assertTrue(os.path.exists(path))

        # After exiting context, should be destroyed
        self.assertFalse(os.path.exists(path))

    def test_context_manager_auto_destroy_false(self):
        """Test that cpuset persists when auto_destroy=False"""
        path = None

        with cpuset.Cpuset('test_context_no_auto', auto_destroy=False) as cs:
            path = cs.cpuset_path
            cs.write_memnode('0')
            cs.assign_cpus('0')
            self.assertTrue(os.path.exists(path))

        # After exiting context, should still exist
        self.assertTrue(os.path.exists(path))

        # Clean up manually
        os.rmdir(path)


@unittest.skipUnless(os.geteuid() == 0, "Requires root permissions")
@unittest.skipUnless(cpuset.CpusetsInit().supported, "Requires cgroup v2 support")
class TestTaskMigration(unittest.TestCase):
    """Test task migration between cpusets"""

    def setUp(self):
        """Set up test fixtures"""
        self.test_cpusets = []
        self.ci = cpuset.CpusetsInit()

    def tearDown(self):
        """Clean up any cpusets created during tests"""
        for cs in self.test_cpusets:
            try:
                if hasattr(cs, 'cpuset_path') and os.path.exists(cs.cpuset_path):
                    tm = cpuset.TaskMigrate(cs, self.ci)
                    tm.migrate()
                    cs.destroy()
            except:
                pass

    def test_write_pid(self):
        """Test writing a PID to a cpuset"""
        cs = cpuset.Cpuset('test_pid')
        self.test_cpusets.append(cs)

        cs.write_memnode('0')
        cs.assign_cpus('0')

        # Try to write our own PID
        pid = os.getpid()
        result = cs.write_pid(pid)

        # Should succeed (we can move ourselves)
        self.assertTrue(result)

        # Verify PID is in the cpuset
        tasks = cs.get_tasks()
        self.assertIn(str(pid), tasks)

        # Move back to root
        self.ci.write_pid(pid)

    def test_task_migrate(self):
        """Test TaskMigrate class"""
        cs = cpuset.Cpuset('test_migrate')
        self.test_cpusets.append(cs)

        cs.write_memnode('0')
        cs.assign_cpus('0')

        # Move our PID to the cpuset
        pid = os.getpid()
        cs.write_pid(pid)

        # Migrate back to root
        tm = cpuset.TaskMigrate(cs, self.ci)
        migrated, failed = tm.migrate()

        self.assertGreaterEqual(migrated, 1)  # At least our process
        self.assertGreaterEqual(failed, 0)  # Some failures are OK


@unittest.skipUnless(os.geteuid() == 0, "Requires root permissions")
@unittest.skipUnless(cpuset.CpusetsInit().supported, "Requires cgroup v2 support")
class TestDiscoveryFunctions(unittest.TestCase):
    """Test tuna-specific discovery functions"""

    def setUp(self):
        """Set up test fixtures"""
        self.test_cpusets = []
        self.ci = cpuset.CpusetsInit()

    def tearDown(self):
        """Clean up any cpusets created during tests"""
        for cs in self.test_cpusets:
            try:
                if hasattr(cs, 'cpuset_path') and os.path.exists(cs.cpuset_path):
                    # Use recursive destroy to handle any nested cpusets
                    cpuset.destroy_cpuset_recursive(cs.cpuset_path, force=True)
            except:
                pass

        # Clean up specific test cpusets that might have nested children
        for name in ['test_parent', 'test_list_1', 'test_list_2',
                     'test_pattern_1', 'test_pattern_2', 'other_cpuset']:
            try:
                cpuset.destroy_cpuset_recursive(name, force=True)
            except:
                pass

    def test_list_cpusets_no_pattern(self):
        """Test listing all cpusets without pattern"""
        # Create some test cpusets
        cs1 = cpuset.Cpuset('test_list_1')
        cs2 = cpuset.Cpuset('test_list_2')
        self.test_cpusets.extend([cs1, cs2])

        cs1.write_memnode('0')
        cs2.write_memnode('0')

        # List all cpusets
        all_cpusets = cpuset.list_cpusets(recursive=False)

        # Should include our test cpusets
        self.assertIn('test_list_1', all_cpusets)
        self.assertIn('test_list_2', all_cpusets)

    def test_list_cpusets_with_pattern(self):
        """Test listing cpusets with pattern matching"""
        # Create test cpusets with different names
        cs1 = cpuset.Cpuset('test_pattern_1')
        cs2 = cpuset.Cpuset('test_pattern_2')
        cs3 = cpuset.Cpuset('other_cpuset')
        self.test_cpusets.extend([cs1, cs2, cs3])

        cs1.write_memnode('0')
        cs2.write_memnode('0')
        cs3.write_memnode('0')

        # List with pattern
        pattern_cpusets = cpuset.list_cpusets(pattern='test_pattern_*', recursive=False)

        # Should only include matching cpusets
        self.assertIn('test_pattern_1', pattern_cpusets)
        self.assertIn('test_pattern_2', pattern_cpusets)
        self.assertNotIn('other_cpuset', pattern_cpusets)

    def test_list_cpusets_recursive(self):
        """Test recursive cpuset listing"""
        # Clean up first in case of previous failed run
        cpuset.destroy_cpuset_recursive('test_parent', force=True)

        # Create nested cpusets using Cpuset class
        # Note: In cgroup v2, creating nested cpusets requires proper setup
        parent = cpuset.Cpuset('test_parent')
        self.test_cpusets.append(parent)
        parent.write_memnode('0')
        parent.assign_cpus('0')

        # Create child as a proper cpuset (not just mkdir)
        # First enable cpuset controller in parent
        try:
            with open(os.path.join(parent.cpuset_path, 'cgroup.subtree_control'), 'w') as f:
                f.write('+cpuset')
        except (OSError, PermissionError):
            pass

        # Create child directory and enable it as a cpuset
        child_path = os.path.join(parent.cpuset_path, 'test_child')
        if not os.path.exists(child_path):
            os.mkdir(child_path)
            # Write to cpuset.cpus to make it a valid cpuset
            try:
                with open(os.path.join(child_path, 'cpuset.cpus'), 'w') as f:
                    f.write('0')
                with open(os.path.join(child_path, 'cpuset.mems'), 'w') as f:
                    f.write('0')
            except (OSError, PermissionError):
                pass

        # List recursively
        all_cpusets = cpuset.list_cpusets(recursive=True)

        # Should find both parent and child
        self.assertIn('test_parent', all_cpusets)
        # Child should be found if it's a valid cpuset
        has_child = os.path.exists(os.path.join(child_path, 'cpuset.cpus'))
        if has_child:
            self.assertIn('test_parent/test_child', all_cpusets)

        # Clean up child first (must destroy children before parents)
        if os.path.exists(child_path):
            try:
                os.rmdir(child_path)
            except OSError:
                pass  # May have processes, tearDown will handle it


@unittest.skipUnless(os.geteuid() == 0, "Requires root permissions")
@unittest.skipUnless(cpuset.CpusetsInit().supported, "Requires cgroup v2 support")
class TestCleanupFunctions(unittest.TestCase):
    """Test tuna-specific cleanup functions"""

    def setUp(self):
        """Set up test fixtures"""
        self.test_cpusets = []
        self.ci = cpuset.CpusetsInit()

    def tearDown(self):
        """Clean up any cpusets that might still exist"""
        # Clean up any test_cleanup_* cpusets (including nested ones)
        try:
            # First try to clean with recursive pattern matching
            cpuset.cleanup_cpusets('test_cleanup_*', force=True, recursive=True)
        except:
            pass

        # Also clean up specific problem cpusets from failed runs
        for name in ['test_cleanup_nested_parent', 'test_cleanup_bulk_1',
                     'test_cleanup_bulk_2', 'other_bulk']:
            try:
                cpuset.destroy_cpuset_recursive(name, force=True)
            except:
                pass

    def test_destroy_cpuset_by_name(self):
        """Test destroying cpuset by name"""
        cs = cpuset.Cpuset('test_cleanup_destroy')
        cs.write_memnode('0')
        path = cs.cpuset_path

        self.assertTrue(os.path.exists(path))

        # Destroy by name
        cpuset.destroy_cpuset('test_cleanup_destroy', force=False)

        self.assertFalse(os.path.exists(path))

    def test_destroy_cpuset_by_path(self):
        """Test destroying cpuset by full path"""
        cs = cpuset.Cpuset('test_cleanup_path')
        cs.write_memnode('0')
        path = cs.cpuset_path

        self.assertTrue(os.path.exists(path))

        # Destroy by full path
        cpuset.destroy_cpuset(path, force=False)

        self.assertFalse(os.path.exists(path))

    def test_destroy_cpuset_with_force(self):
        """Test destroying cpuset with force (migrates tasks first)"""
        cs = cpuset.Cpuset('test_cleanup_force')
        cs.write_memnode('0')
        cs.assign_cpus('0')

        # Move our PID to the cpuset
        cs.write_pid(os.getpid())

        # Destroy with force (should migrate tasks first)
        cpuset.destroy_cpuset('test_cleanup_force', force=True)

        self.assertFalse(os.path.exists(cs.cpuset_path))

        # Verify we're back in root cgroup
        with open('/proc/self/cgroup') as f:
            cgroup_info = f.read()
        # In cgroup v2, processes in root are at '0::/'
        self.assertIn('0::/', cgroup_info)

    def test_destroy_nonexistent_cpuset_raises_error(self):
        """Test that destroying non-existent cpuset raises OSError"""
        with self.assertRaises(OSError):
            cpuset.destroy_cpuset('nonexistent_cpuset')

    def test_cleanup_cpusets_with_pattern(self):
        """Test cleanup_cpusets with pattern matching"""
        # Create multiple cpusets
        cs1 = cpuset.Cpuset('test_cleanup_bulk_1')
        cs2 = cpuset.Cpuset('test_cleanup_bulk_2')
        cs3 = cpuset.Cpuset('other_bulk')

        cs1.write_memnode('0')
        cs2.write_memnode('0')
        cs3.write_memnode('0')

        # Cleanup with pattern
        cpuset.cleanup_cpusets('test_cleanup_bulk_*', force=True, recursive=False)

        # Pattern-matched cpusets should be gone
        self.assertFalse(os.path.exists(cs1.cpuset_path))
        self.assertFalse(os.path.exists(cs2.cpuset_path))

        # Other cpuset should still exist
        self.assertTrue(os.path.exists(cs3.cpuset_path))

        # Clean up remaining
        cs3.destroy()

    def test_cleanup_nested_cpusets(self):
        """Test cleanup destroys children before parents"""
        # Clean up first in case of previous failed run
        cpuset.destroy_cpuset_recursive('test_cleanup_nested_parent', force=True)

        # Create parent
        parent = cpuset.Cpuset('test_cleanup_nested_parent')
        parent.write_memnode('0')
        parent.assign_cpus('0')

        # Enable cpuset controller in parent for children
        try:
            with open(os.path.join(parent.cpuset_path, 'cgroup.subtree_control'), 'w') as f:
                f.write('+cpuset')
        except (OSError, PermissionError):
            pass

        # Create children as proper cpusets
        child1_path = os.path.join(parent.cpuset_path, 'child1')
        child2_path = os.path.join(parent.cpuset_path, 'child2')

        for child_path in [child1_path, child2_path]:
            if not os.path.exists(child_path):
                os.mkdir(child_path)
                # Make it a valid cpuset
                try:
                    with open(os.path.join(child_path, 'cpuset.cpus'), 'w') as f:
                        f.write('0')
                    with open(os.path.join(child_path, 'cpuset.mems'), 'w') as f:
                        f.write('0')
                except (OSError, PermissionError):
                    pass

        # Cleanup recursively
        cpuset.cleanup_cpusets('test_cleanup_nested_*', force=False, recursive=True)

        # Everything should be gone
        self.assertFalse(os.path.exists(child1_path))
        self.assertFalse(os.path.exists(child2_path))
        self.assertFalse(os.path.exists(parent.cpuset_path))


class TestProcessMovement(unittest.TestCase):
    """Test moving individual processes to cpusets"""

    @unittest.skipUnless(os.geteuid() == 0, "Requires root")
    def test_write_pid_self(self):
        """Test moving current process to a cpuset"""
        # Pre-cleanup
        cpuset.cleanup_cpusets('test_write_pid*', force=True, recursive=True)

        cs = cpuset.Cpuset('test_write_pid')
        try:
            # Configure cpuset
            cs.write_memnode('0')
            cs.assign_cpus('0-1')

            # Move current process
            current_pid = os.getpid()
            result = cs.write_pid(current_pid)
            self.assertTrue(result, "write_pid should return True on success")

            # Verify it's in the cpuset
            tasks = cs.get_tasks()
            self.assertIn(str(current_pid), tasks, "Current PID should be in cpuset")

            # Move back to root
            root = cpuset.CpusetsInit()
            root.write_pid(current_pid)
        finally:
            # Cleanup
            cpuset.destroy_cpuset('test_write_pid', force=True)

    @unittest.skipUnless(os.geteuid() == 0, "Requires root")
    def test_write_pid_invalid(self):
        """Test write_pid with invalid PID"""
        # Pre-cleanup
        cpuset.cleanup_cpusets('test_write_invalid*', force=True, recursive=True)

        cs = cpuset.Cpuset('test_write_invalid')
        try:
            cs.write_memnode('0')
            cs.assign_cpus('0-1')

            # Try to move non-existent PID
            result = cs.write_pid(999999)
            self.assertFalse(result, "write_pid should return False for invalid PID")
        finally:
            # Cleanup
            cpuset.destroy_cpuset('test_write_invalid', force=True)

    @unittest.skipUnless(os.geteuid() == 0, "Requires root")
    def test_multiple_pids(self):
        """Test moving multiple processes to a cpuset"""
        # Pre-cleanup
        cpuset.cleanup_cpusets('test_multi_pid*', force=True, recursive=True)

        cs = cpuset.Cpuset('test_multi_pid')
        try:
            cs.write_memnode('0')
            cs.assign_cpus('0-1')

            # Move current process
            current_pid = os.getpid()
            cs.write_pid(current_pid)

            # Get current tasks
            tasks = cs.get_tasks()
            self.assertIn(str(current_pid), tasks)

            # Move back to root
            root = cpuset.CpusetsInit()
            root.write_pid(current_pid)
        finally:
            cpuset.destroy_cpuset('test_multi_pid', force=True)


class TestNumaDetection(unittest.TestCase):
    """Test NUMA topology detection functions"""

    def test_discover_numa_nodes(self):
        """Test that NUMA node discovery returns valid data"""
        nodes = cpuset.discover_numa_nodes()

        # Should have at least one node
        self.assertGreaterEqual(len(nodes), 1)

        # All nodes should have node IDs starting from 0
        self.assertIn(0, nodes)

        # Each node should have CPUs
        for nodeid, node in nodes.items():
            self.assertIsInstance(nodeid, int)
            self.assertIsInstance(node, cpuset.NumaNode)
            self.assertEqual(node.nodeid, nodeid)
            self.assertIsInstance(node.cpus, set)

    def test_get_numa_nodes_for_cpus_basic(self):
        """Test NUMA node detection for basic CPU lists"""
        # Get available nodes
        nodes = cpuset.discover_numa_nodes()

        # Get some CPUs from first node
        first_node = nodes[min(nodes.keys())]
        test_cpus = list(first_node.cpus)[:4] if len(first_node.cpus) >= 4 else list(first_node.cpus)

        if test_cpus:
            result = cpuset.get_numa_nodes_for_cpus(test_cpus)
            # Should return a string
            self.assertIsInstance(result, str)
            # Should contain at least the first node
            self.assertIn(str(min(nodes.keys())), result)

    def test_get_numa_nodes_for_empty_list(self):
        """Test NUMA node detection for empty CPU list"""
        result = cpuset.get_numa_nodes_for_cpus([])
        # Should default to "0"
        self.assertEqual(result, "0")

    def test_get_numa_nodes_for_single_cpu(self):
        """Test NUMA node detection for single CPU"""
        # Get first available CPU
        nodes = cpuset.discover_numa_nodes()
        first_node = nodes[min(nodes.keys())]

        if first_node.cpus:
            test_cpu = list(first_node.cpus)[0]
            result = cpuset.get_numa_nodes_for_cpus([test_cpu])
            # Should return a valid node string
            self.assertIsInstance(result, str)
            self.assertTrue(len(result) > 0)

    def test_numa_node_contains(self):
        """Test NumaNode __contains__ operator"""
        node = cpuset.NumaNode(0, {0, 1, 2, 3})
        self.assertIn(0, node)
        self.assertIn(3, node)
        self.assertNotIn(4, node)
        self.assertNotIn(100, node)

    def test_numa_node_int(self):
        """Test NumaNode __int__ operator"""
        node = cpuset.NumaNode(5, {0, 1})
        self.assertEqual(int(node), 5)


def suite():
    """Create test suite"""
    loader = unittest.TestLoader()
    suite = unittest.TestSuite()

    # Add all test classes
    suite.addTests(loader.loadTestsFromTestCase(TestCpusetsInit))
    suite.addTests(loader.loadTestsFromTestCase(TestNumaDetection))
    suite.addTests(loader.loadTestsFromTestCase(TestCpusetCreation))
    suite.addTests(loader.loadTestsFromTestCase(TestCpusetConfiguration))
    suite.addTests(loader.loadTestsFromTestCase(TestContextManager))
    suite.addTests(loader.loadTestsFromTestCase(TestTaskMigration))
    suite.addTests(loader.loadTestsFromTestCase(TestDiscoveryFunctions))
    suite.addTests(loader.loadTestsFromTestCase(TestCleanupFunctions))
    suite.addTests(loader.loadTestsFromTestCase(TestProcessMovement))

    return suite


if __name__ == '__main__':
    # Check if running as root for tests that need it
    if os.geteuid() != 0:
        print("\n" + "="*60)
        print("WARNING: Not running as root")
        print("="*60)
        print("Most tests require root permissions to create/destroy cgroups.")
        print("Only basic detection tests will run.")
        print("\nTo run all tests:")
        print("  sudo python3 -m unittest tests.test_cpuset")
        print("="*60 + "\n")

    # Run tests
    runner = unittest.TextTestRunner(verbosity=2)
    result = runner.run(suite())

    # Exit with appropriate code
    sys.exit(0 if result.wasSuccessful() else 1)
