NetAlertX/test/test_device_atomicity.py

"""
Test for atomicity of device updates with source-tracking.

Verifies that:
1. If source-tracking fails, the device row is rolled back.
2. If source-tracking succeeds, device row and sources are both committed.
3. Database remains consistent in both scenarios.
"""

import sys
import os
import sqlite3
import tempfile
import unittest
from unittest.mock import patch


# Add server and plugins to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'server'))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'server', 'plugins'))
sys.path.insert(0, os.path.dirname(__file__))

from models.device_instance import DeviceInstance  # noqa: E402 [flake8 lint suppression]
from plugin_helper import normalize_mac  # noqa: E402 [flake8 lint suppression]
from db_test_helpers import CREATE_DEVICES  # noqa: E402 [flake8 lint suppression]


class TestDeviceAtomicity(unittest.TestCase):
    """Test atomic transactions for device updates with source-tracking."""

    def setUp(self):
        """Create an in-memory SQLite DB for testing."""
        self.test_db = tempfile.NamedTemporaryFile(delete=False, suffix='.db')
        self.test_db_path = self.test_db.name
        self.test_db.close()

        # Create minimal schema
        conn = sqlite3.connect(self.test_db_path)
        conn.row_factory = sqlite3.Row
        cur = conn.cursor()

        # Create Devices table with source-tracking columns - reuses the
        # canonical schema shared by the rest of the suite rather than a
        # local copy (the local copy this replaced had drifted: it declared
        # devTypeSource/devIconSource/devGroupSource/devLocationSource/
        # devCommentsSource, none of which exist in
        # server/db/schema/app.sql or are read/written by any production
        # code - source tracking only exists for the fields CREATE_DEVICES
        # actually declares, e.g. devNameSource/devVendorSource/devMacSource).
        cur.execute(CREATE_DEVICES)
        conn.commit()
        conn.close()

    def tearDown(self):
        """Clean up test database."""
        if os.path.exists(self.test_db_path):
            os.unlink(self.test_db_path)

    def _get_test_db_connection(self):
        """Override database connection for testing."""
        conn = sqlite3.connect(self.test_db_path)
        conn.row_factory = sqlite3.Row
        return conn

    def test_create_new_device_atomicity(self):
        """
        Test that device creation and source-tracking are atomic.
        If source tracking fails, the device should not be created.
        """
        device_instance = DeviceInstance()
        test_mac = normalize_mac("aa:bb:cc:dd:ee:ff")

        # Patch at module level where it's used
        with patch('models.device_instance.get_temp_db_connection', self._get_test_db_connection):
            # Create a new device
            data = {
                "createNew": True,
                "devMac": test_mac,
                "devName": "Test Device",
                "devOwner": "John Doe",
                "devType": "Laptop",
            }

            result = device_instance.setDeviceData(test_mac, data)

            # Verify success
            self.assertTrue(result["success"], f"Device creation failed: {result}")

            # Verify device exists
            conn = self._get_test_db_connection()
            cur = conn.cursor()
            cur.execute("SELECT * FROM Devices WHERE devMac = ?", (test_mac,))
            device = cur.fetchone()
            conn.close()

            self.assertIsNotNone(device, "Device was not created")
            self.assertEqual(device["devName"], "Test Device")

            # Verify source tracking was set
            self.assertEqual(device["devMacSource"], "NEWDEV")
            self.assertEqual(device["devNameSource"], "NEWDEV")

    def test_update_device_with_source_tracking_atomicity(self):
        """
        Test that device update and source-tracking are atomic.
        If source tracking fails, the device update should be rolled back.
        """
        device_instance = DeviceInstance()
        test_mac = normalize_mac("aa:bb:cc:dd:ee:ff")

        # Create initial device
        conn = self._get_test_db_connection()
        cur = conn.cursor()
        cur.execute("""
            INSERT INTO Devices (
                devMac, devName, devOwner, devType, devNameSource
            ) VALUES (?, ?, ?, ?, ?)
                    """, (test_mac, "Old Name", "Old Owner", "Desktop", "PLUGIN"))
        conn.commit()
        conn.close()

        # Patch database connection
        with patch('models.device_instance.get_temp_db_connection', self._get_test_db_connection):
            with patch('models.device_instance.enforce_source_on_user_update') as mock_enforce:
                mock_enforce.return_value = None
                data = {
                    "createNew": False,
                    "devMac": test_mac,
                    "devName": "New Name",
                    "devOwner": "New Owner",
                }

                result = device_instance.setDeviceData(test_mac, data)

            # Verify success
            self.assertTrue(result["success"], f"Device update failed: {result}")

            # Verify device was updated
            conn = self._get_test_db_connection()
            cur = conn.cursor()
            cur.execute("SELECT * FROM Devices WHERE devMac = ?", (test_mac,))
            device = cur.fetchone()
            conn.close()

            self.assertEqual(device["devName"], "New Name")
            self.assertEqual(device["devOwner"], "New Owner")

    def test_source_tracking_failure_rolls_back_device(self):
        """
        Test that if enforce_source_on_user_update fails, the entire
        transaction is rolled back (device and sources).
        """
        device_instance = DeviceInstance()
        test_mac = normalize_mac("aa:bb:cc:dd:ee:ff")

        # Create initial device
        conn = self._get_test_db_connection()
        cur = conn.cursor()
        cur.execute("""
            INSERT INTO Devices (
                devMac, devName, devOwner, devType, devNameSource
            ) VALUES (?, ?, ?, ?, ?)
                    """, (test_mac, "Original Name", "Original Owner", "Desktop", "PLUGIN"))
        conn.commit()
        conn.close()

        # Patch database connection and mock source enforcement failure
        with patch('models.device_instance.get_temp_db_connection', self._get_test_db_connection):
            with patch('models.device_instance.enforce_source_on_user_update') as mock_enforce:
                # Simulate source tracking failure
                mock_enforce.side_effect = Exception("Source tracking error")

                data = {
                    "createNew": False,
                    "devMac": test_mac,
                    "devName": "Failed Update",
                    "devOwner": "Failed Owner",
                }

                result = device_instance.setDeviceData(test_mac, data)

            # Verify error response
            self.assertFalse(result["success"])
            self.assertIn("Source tracking failed", result["error"])

            # Verify device was NOT updated (rollback successful)
            conn = self._get_test_db_connection()
            cur = conn.cursor()
            cur.execute("SELECT * FROM Devices WHERE devMac = ?", (test_mac,))
            device = cur.fetchone()
            conn.close()

            self.assertEqual(device["devName"], "Original Name", "Device should not have been updated on source tracking failure")
            self.assertEqual(device["devOwner"], "Original Owner", "Device should not have been updated on source tracking failure")


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