diff --git a/netbox_custom_objects/__init__.py b/netbox_custom_objects/__init__.py index af41097..f1586b9 100644 --- a/netbox_custom_objects/__init__.py +++ b/netbox_custom_objects/__init__.py @@ -32,6 +32,42 @@ def _migration_finished(sender, **kwargs): _migrations_checked = None +# Module-level flag so the heal runs at most once per process invocation even +# though post_migrate fires once per installed app. +_heal_ran = False + + +def _heal_mixin_columns(sender, **kwargs): + """ + post_migrate signal handler: detect and apply mixin column drift. + + Fires after every 'manage.py migrate' run (once per installed app). The + module-level _heal_ran flag ensures the actual work happens only once per + process so the cost is negligible on normal server starts where no + migrations run. + + Skipped during makemigrations and collectstatic (DB may be unavailable or + in an inconsistent state for our purposes). + """ + global _heal_ran + if _heal_ran: + return + + if any(cmd in sys.argv for cmd in ("makemigrations", "collectstatic")): + return + + _heal_ran = True + + try: + from netbox_custom_objects.mixin_migration import heal_all_cots # noqa: PLC0415 + heal_all_cots(verbosity=kwargs.get("verbosity", 1)) + except Exception: + import logging # noqa: PLC0415 + logging.getLogger(__name__).exception( + "upgrade_custom_objects: unexpected error during mixin drift check" + ) + + def _patch_object_selector_view(): """ Patch ObjectSelectorView to support dynamically-generated custom object models. @@ -180,6 +216,9 @@ def ready(self): pre_migrate.connect(_migration_started) post_migrate.connect(_migration_finished) + # Heal mixin column drift after every migrate run (issue #391 Phase 2) + post_migrate.connect(_heal_mixin_columns) + # Patch ObjectSelectorView to support dynamically-generated custom object models _patch_object_selector_view() diff --git a/netbox_custom_objects/management/__init__.py b/netbox_custom_objects/management/__init__.py new file mode 100644 index 0000000..447ca19 --- /dev/null +++ b/netbox_custom_objects/management/__init__.py @@ -0,0 +1 @@ +# Management commands for netbox_custom_objects diff --git a/netbox_custom_objects/management/commands/__init__.py b/netbox_custom_objects/management/commands/__init__.py new file mode 100644 index 0000000..447ca19 --- /dev/null +++ b/netbox_custom_objects/management/commands/__init__.py @@ -0,0 +1 @@ +# Management commands for netbox_custom_objects diff --git a/netbox_custom_objects/management/commands/upgrade_custom_objects.py b/netbox_custom_objects/management/commands/upgrade_custom_objects.py new file mode 100644 index 0000000..21fbf05 --- /dev/null +++ b/netbox_custom_objects/management/commands/upgrade_custom_objects.py @@ -0,0 +1,126 @@ +""" +management command: upgrade_custom_objects + +Checks all Custom Object Type tables for mixin column drift and applies safe +fixes. Intended as an explicit escape hatch alongside the automatic +post_migrate signal handler (issue #391). + +Usage examples +-------------- + # Check and fix all COTs + manage.py upgrade_custom_objects + + # Preview changes without touching the DB + manage.py upgrade_custom_objects --dry-run + + # Operate on a single COT (by name or numeric ID) + manage.py upgrade_custom_objects --cot my_device + manage.py upgrade_custom_objects --cot 7 --dry-run +""" + +from django.core.management.base import BaseCommand, CommandError + +from netbox_custom_objects.mixin_migration import heal_cot + + +class Command(BaseCommand): + help = ( + "Detect and apply mixin column drift for Custom Object Type tables. " + "New columns contributed by the CustomObject base class (e.g. from a " + "NetBox upgrade) are added automatically when nullable or defaulted. " + "Non-nullable columns without defaults and column removals are reported " + "but never applied automatically." + ) + + def add_arguments(self, parser): + parser.add_argument( + "--dry-run", + action="store_true", + help="Report what would change without making any DB modifications.", + ) + parser.add_argument( + "--cot", + metavar="NAME_OR_ID", + help="Limit to a single Custom Object Type (name or numeric ID).", + ) + + def handle(self, *args, **options): + dry_run = options["dry_run"] + cot_filter = options.get("cot") + verbosity = options["verbosity"] + + if dry_run: + self.stdout.write(self.style.WARNING("DRY RUN — no changes will be made.\n")) + + if cot_filter: + from netbox_custom_objects.models import CustomObjectType # noqa: PLC0415 + try: + if cot_filter.isdigit(): + cot = CustomObjectType.objects.get(pk=int(cot_filter)) + else: + cot = CustomObjectType.objects.get(name=cot_filter) + except CustomObjectType.DoesNotExist: + raise CommandError(f"No Custom Object Type found: {cot_filter!r}") + + result = heal_cot(cot, verbosity=verbosity, dry_run=dry_run) + self._print_cot_result(cot.name, result, dry_run) + else: + from netbox_custom_objects.models import CustomObjectType # noqa: PLC0415 + cots = list(CustomObjectType.objects.all()) + total = len(cots) + healed = warnings = 0 + for cot in cots: + result = heal_cot(cot, verbosity=verbosity, dry_run=dry_run) + self._print_cot_result(cot.name, result, dry_run) + if result["added"]: + healed += 1 + warnings += len(result["warned"]) + self._print_summary( + {"total": total, "healed": healed, "warnings": warnings}, + dry_run, + ) + + # ------------------------------------------------------------------ + # Output helpers + # ------------------------------------------------------------------ + + def _print_cot_result(self, cot_name, result, dry_run): + added = result["added"] + warned = result["warned"] + + if not added and not warned: + self.stdout.write( + self.style.SUCCESS(f"COT {cot_name!r}: no drift detected.") + ) + return + + tag = " [DRY RUN]" if dry_run else "" + for field_name in added: + self.stdout.write( + self.style.SUCCESS(f" {tag} + Added column: {field_name}") + ) + for entry in warned: + self.stdout.write( + self.style.WARNING(f" ! {entry['message']}") + ) + + def _print_summary(self, summary, dry_run): + tag = " (dry run)" if dry_run else "" + if summary["healed"] == 0 and summary["warnings"] == 0: + self.stdout.write( + self.style.SUCCESS( + f"All {summary['total']} COT table(s) are up to date{tag}." + ) + ) + else: + self.stdout.write( + f"{summary['total']} COT(s) checked{tag}: " + f"{summary['healed']} healed, " + f"{summary['warnings']} warning(s)." + ) + if summary["warnings"]: + self.stdout.write( + self.style.WARNING( + "Run with -v 2 or check the application log for warning details." + ) + ) diff --git a/netbox_custom_objects/mixin_migration.py b/netbox_custom_objects/mixin_migration.py new file mode 100644 index 0000000..757232e --- /dev/null +++ b/netbox_custom_objects/mixin_migration.py @@ -0,0 +1,274 @@ +""" +Mixin column drift detection and repair for Custom Object Type tables. + +Phase 2 of issue #391: when NetBox is upgraded and a mixin (e.g. +ChangeLoggingMixin) gains a new concrete column, existing COT tables will be +missing that column. This module provides: + + heal_cot(cot, verbosity, dry_run) — check and repair a single COT table + heal_all_cots(verbosity, dry_run) — iterate over all COTs + +Both are called from: + - The post_migrate signal handler in __init__.py (automatic, zero-config) + - The upgrade_custom_objects management command (explicit, with --dry-run) + +Safety rules +------------ + ADD allowed : new column is nullable OR has a Django-level default + Warn only : new column is NOT NULL with no default (would fail for existing rows) + Warn only : column type appears to have changed + Never : auto-drop a column that is no longer in the base class +""" + +import logging + +from django.db import connection + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + +def _expected_base_fields(cot, model=None): + """ + Return {db_column_name: Django field instance} for every concrete column + that the current CustomObject mixin hierarchy contributes to *cot*'s DB + table, excluding user-defined fields. + + Keyed by f.column (the actual DB column name) so results can be compared + directly against _actual_column_names() output, which returns DB column + names from introspection. Using f.name would produce incorrect comparisons + for FK fields (where f.name='foo' but f.column='foo_id') or any field that + overrides db_column. + + User fields are excluded by matching against their Python attribute names + (f.name). This is equivalent to matching by f.column for user-defined COT + fields because they are never created with db_column overrides. + + Pass *model* to avoid a second get_model() call when the caller already + holds the model reference. + """ + if model is None: + model = cot.get_model() + user_field_names = set(cot.fields.values_list("name", flat=True)) + return { + f.column: f + for f in model._meta.concrete_fields + if f.name not in user_field_names + } + + +def _actual_column_names(table_name): + """ + Return the set of column names currently present in *table_name*. + + Raises OperationalError / ProgrammingError if the table does not exist. + """ + with connection.cursor() as cursor: + return { + col.name + for col in connection.introspection.get_table_description(cursor, table_name) + } + + +def _can_auto_add(field): + """ + Return True if it is safe to ADD COLUMN for *field* on a table that + already has rows. + + A column is safe to add when existing rows can receive a value without + violating constraints: + - Nullable columns default to NULL for existing rows. + - Columns with a Django-level default use that value. + """ + return field.null or field.has_default() + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + +def heal_cot(cot, verbosity=1, dry_run=False): + """ + Detect and repair mixin column drift for a single CustomObjectType. + + Parameters + ---------- + cot : CustomObjectType instance + verbosity : int 0=silent, 1=changes+warnings, 2=verbose + dry_run : bool if True, report but do not modify the DB + + Returns + ------- + dict with keys: + "added" : list of column names successfully added (or would-be added) + "warned" : list of dicts {type, field, message} for non-auto-fixable issues + """ + table_name = cot.get_database_table_name() + added = [] + warned = [] + + try: + actual_names = _actual_column_names(table_name) + except Exception as exc: + logger.warning( + "upgrade_custom_objects: cannot introspect table %r (COT %s): %s", + table_name, cot.pk, exc, + ) + return {"added": added, "warned": warned} + + # Resolve model once; pass it through to avoid a duplicate get_model() call. + model = cot.get_model() + expected = _expected_base_fields(cot, model) + + # Build a lookup of what was stored in the last snapshot for type comparison. + # Note: schema_document["base_columns"] stores column names as f.name (from + # Phase 1's _collect_base_columns). For all current base fields f.name == + # f.column, so this lookup is consistent with expected's f.column keys. + stored_col_info = { + c["name"]: c + for c in (cot.schema_document or {}).get("base_columns", []) + } + + # ── New columns in expected but missing from actual ────────────────────── + for col_name, field in expected.items(): + if col_name in actual_names: + continue + + if not _can_auto_add(field): + entry = { + "type": "new_non_nullable", + "field": col_name, + "message": ( + f"Table {table_name!r}: new base column {col_name!r} " + f"({field.__class__.__name__}) is NOT NULL with no default — " + f"cannot be added automatically. Add a default or make it " + f"nullable upstream, then re-run 'manage.py upgrade_custom_objects'." + ), + } + warned.append(entry) + logger.warning(entry["message"]) + continue + + if dry_run: + added.append(col_name) + continue + + try: + with connection.schema_editor() as editor: + editor.add_field(model, field) + added.append(col_name) + if verbosity >= 1: + logger.info( + "upgrade_custom_objects: added column %r to table %r", + col_name, table_name, + ) + except Exception as exc: + entry = { + "type": "add_failed", + "field": col_name, + "message": ( + f"Failed to ADD COLUMN {col_name!r} to {table_name!r}: {exc}" + ), + } + warned.append(entry) + logger.error(entry["message"]) + + # ── Type changes on columns present in both expected and actual ────────── + for col_name, field in expected.items(): + if col_name not in actual_names: + continue # already handled above as a new column + stored = stored_col_info.get(col_name) + if not stored or not stored.get("field_class"): + continue # no prior snapshot to compare against + if stored["field_class"] != field.__class__.__name__: + entry = { + "type": "type_changed", + "field": col_name, + "message": ( + f"Table {table_name!r}: column {col_name!r} type may have changed " + f"(was {stored['field_class']!r}, now {field.__class__.__name__!r}). " + f"Manual inspection and migration required." + ), + } + warned.append(entry) + logger.warning(entry["message"]) + + # ── Columns removed from base class but still in DB ───────────────────── + stored_base_names = set(stored_col_info) + for col_name in sorted(stored_base_names - set(expected)): + if col_name in actual_names: + entry = { + "type": "removed_from_model", + "field": col_name, + "message": ( + f"Table {table_name!r}: column {col_name!r} still exists in the " + f"database but is no longer in the CustomObject base class. " + f"Manual cleanup may be required." + ), + } + warned.append(entry) + logger.warning(entry["message"]) + + # ── Refresh snapshot after successful additions ────────────────────────── + if added and not dry_run: + # We cannot use _store_base_column_snapshot(model) here because the + # generated model's _meta is built from the CustomObject class definition + # and does not include columns added directly to the DB by this heal pass. + # Instead, merge the newly-added field info into the existing snapshot. + doc = cot.schema_document or {} + current_cols = {c["name"]: c for c in doc.get("base_columns", [])} + for col_name in added: + field = expected[col_name] + current_cols[col_name] = { + "name": col_name, + "field_class": field.__class__.__name__, + "null": field.null, + } + doc["base_columns"] = list(current_cols.values()) + cot.__class__.objects.filter(pk=cot.pk).update(schema_document=doc) + cot.schema_document = doc + + return {"added": added, "warned": warned} + + +def heal_all_cots(verbosity=1, dry_run=False): + """ + Run heal_cot() for every CustomObjectType. + + Called by the post_migrate signal handler. The upgrade_custom_objects + management command iterates COTs directly so it can print per-COT output + to stdout. + + Returns + ------- + dict with keys: + "total" : number of COTs checked + "healed" : number of COTs that had columns added + "warnings" : total number of non-auto-fixable issues + """ + from netbox_custom_objects.models import CustomObjectType # noqa: PLC0415 + + total = healed = warnings = 0 + + for cot in CustomObjectType.objects.all(): + total += 1 + result = heal_cot(cot, verbosity=verbosity, dry_run=dry_run) + if result["added"]: + healed += 1 + warnings += len(result["warned"]) + + if verbosity >= 2: + logger.info( + "upgrade_custom_objects: %d COT(s) checked, %d healed, %d warning(s)", + total, healed, warnings, + ) + elif verbosity >= 1 and (healed > 0 or warnings > 0): + logger.info( + "upgrade_custom_objects: %d COT(s) healed, %d warning(s)", + healed, warnings, + ) + + return {"total": total, "healed": healed, "warnings": warnings} diff --git a/netbox_custom_objects/tests/test_mixin_migration.py b/netbox_custom_objects/tests/test_mixin_migration.py new file mode 100644 index 0000000..d02b5bd --- /dev/null +++ b/netbox_custom_objects/tests/test_mixin_migration.py @@ -0,0 +1,392 @@ +""" +Tests for the mixin column drift detection and repair (issue #391, Phase 2). + +Covers: +- _expected_base_fields(): returns the correct base fields, excludes user fields +- _can_auto_add(): correct classification of nullable / defaulted fields +- heal_cot(): detects missing columns, adds safe ones, warns on unsafe ones, + never drops, updates schema_document snapshot after healing, + dry_run mode reports without modifying +- heal_all_cots(): iterates all COTs and returns correct summary counts +- upgrade_custom_objects management command: --dry-run and --cot flags +""" + +from unittest.mock import MagicMock, patch + +from django.db import connection +from django.test import TestCase, TransactionTestCase + +from netbox_custom_objects.mixin_migration import ( + _can_auto_add, + _expected_base_fields, + heal_all_cots, + heal_cot, +) +from netbox_custom_objects.models import CustomObjectType + +from .base import CustomObjectsTestCase, TransactionCleanupMixin + + +# --------------------------------------------------------------------------- +# _can_auto_add() +# --------------------------------------------------------------------------- + +class CanAutoAddTestCase(TestCase): + """Unit tests for _can_auto_add() — no DB required.""" + + def _field(self, null=False, has_default=False, default_value=None): + f = MagicMock() + f.null = null + f.has_default.return_value = has_default + return f + + def test_nullable_field_is_safe(self): + self.assertTrue(_can_auto_add(self._field(null=True))) + + def test_field_with_default_is_safe(self): + self.assertTrue(_can_auto_add(self._field(has_default=True))) + + def test_nullable_and_has_default_is_safe(self): + self.assertTrue(_can_auto_add(self._field(null=True, has_default=True))) + + def test_non_nullable_no_default_is_unsafe(self): + self.assertFalse(_can_auto_add(self._field(null=False, has_default=False))) + + +# --------------------------------------------------------------------------- +# _expected_base_fields() +# --------------------------------------------------------------------------- + +class ExpectedBaseFieldsTestCase( + TransactionCleanupMixin, CustomObjectsTestCase, TransactionTestCase +): + """Tests for _expected_base_fields() — requires a live COT.""" + + def test_returns_id_created_last_updated(self): + cot = self.create_custom_object_type(name="ebf_basic", slug="ebf-basic") + fields = _expected_base_fields(cot) + self.assertIn("id", fields) + self.assertIn("created", fields) + self.assertIn("last_updated", fields) + + def test_excludes_user_defined_field(self): + cot = self.create_custom_object_type(name="ebf_user", slug="ebf-user") + self.create_custom_object_type_field(cot, name="my_col", type="text") + fields = _expected_base_fields(cot) + self.assertNotIn("my_col", fields) + + def test_returns_django_field_instances(self): + from django.db.models import Field + cot = self.create_custom_object_type(name="ebf_inst", slug="ebf-inst") + for field in _expected_base_fields(cot).values(): + self.assertIsInstance(field, Field) + + +# --------------------------------------------------------------------------- +# heal_cot() — normal path +# --------------------------------------------------------------------------- + +class HealCotTestCase( + TransactionCleanupMixin, CustomObjectsTestCase, TransactionTestCase +): + """Integration tests for heal_cot() against a real DB.""" + + def test_no_drift_returns_empty_results(self): + cot = self.create_custom_object_type(name="hc_nodrift", slug="hc-nodrift") + result = heal_cot(cot) + self.assertEqual(result["added"], []) + self.assertEqual(result["warned"], []) + + def test_missing_nullable_column_is_added(self): + """ + Simulate a new nullable base column appearing in the mixin by patching + _expected_base_fields to return an extra field, then verifying that + heal_cot adds it to the actual DB table. + """ + cot = self.create_custom_object_type(name="hc_add", slug="hc-add") + table_name = cot.get_database_table_name() + + # Confirm the column doesn't exist yet + with connection.cursor() as cur: + actual_before = { + c.name for c in connection.introspection.get_table_description(cur, table_name) + } + self.assertNotIn("new_nullable_col", actual_before) + + # Build a real nullable CharField to inject + from django.db import models as dj_models + new_field = dj_models.CharField(max_length=50, null=True, blank=True) + new_field.name = "new_nullable_col" + new_field.column = "new_nullable_col" + new_field.set_attributes_from_name("new_nullable_col") + new_field.model = cot.get_model() + + base_fields = _expected_base_fields(cot) + base_fields["new_nullable_col"] = new_field + + with patch( + "netbox_custom_objects.mixin_migration._expected_base_fields", + return_value=base_fields, + ): + result = heal_cot(cot, verbosity=0) + + self.assertIn("new_nullable_col", result["added"]) + self.assertEqual(result["warned"], []) + + # Verify the column now exists in the DB + with connection.cursor() as cur: + actual_after = { + c.name for c in connection.introspection.get_table_description(cur, table_name) + } + self.assertIn("new_nullable_col", actual_after) + + # Clean up the added column so tearDown can drop the table cleanly + with connection.schema_editor() as editor: + editor.remove_field(cot.get_model(), new_field) + + def test_missing_non_nullable_no_default_produces_warning(self): + """A NOT NULL column without a default cannot be auto-added; must warn.""" + cot = self.create_custom_object_type(name="hc_warn", slug="hc-warn") + + from django.db import models as dj_models + bad_field = dj_models.IntegerField() + bad_field.name = "required_int" + bad_field.column = "required_int" + bad_field.set_attributes_from_name("required_int") + bad_field.model = cot.get_model() + + base_fields = _expected_base_fields(cot) + base_fields["required_int"] = bad_field + + with patch( + "netbox_custom_objects.mixin_migration._expected_base_fields", + return_value=base_fields, + ): + result = heal_cot(cot, verbosity=0) + + self.assertEqual(result["added"], []) + self.assertEqual(len(result["warned"]), 1) + self.assertEqual(result["warned"][0]["type"], "new_non_nullable") + self.assertEqual(result["warned"][0]["field"], "required_int") + + def test_snapshot_updated_after_addition(self): + """schema_document['base_columns'] must be refreshed after columns are added.""" + cot = self.create_custom_object_type(name="hc_snap", slug="hc-snap") + + from django.db import models as dj_models + extra_field = dj_models.CharField(max_length=10, null=True, blank=True) + extra_field.name = "snap_col" + extra_field.column = "snap_col" + extra_field.set_attributes_from_name("snap_col") + extra_field.model = cot.get_model() + + base_fields = _expected_base_fields(cot) + base_fields["snap_col"] = extra_field + + with patch( + "netbox_custom_objects.mixin_migration._expected_base_fields", + return_value=base_fields, + ): + heal_cot(cot, verbosity=0) + + cot.refresh_from_db() + names = {c["name"] for c in cot.schema_document.get("base_columns", [])} + self.assertIn("snap_col", names) + + # Clean up + with connection.schema_editor() as editor: + editor.remove_field(cot.get_model(), extra_field) + + def test_removed_column_produces_warning_not_drop(self): + """A column in schema_document['base_columns'] but removed from model must only warn.""" + cot = self.create_custom_object_type(name="hc_drop", slug="hc-drop") + + # Add ghost_col to the actual DB table so the heal checker sees it. + # This simulates a column that was once part of a mixin but has since + # been removed from the CustomObject base class. + from django.db import models as dj_models + ghost_field = dj_models.CharField(max_length=50, null=True, blank=True) + ghost_field.name = "ghost_col" + ghost_field.column = "ghost_col" + ghost_field.set_attributes_from_name("ghost_col") + ghost_field.model = cot.get_model() + with connection.schema_editor() as editor: + editor.add_field(cot.get_model(), ghost_field) + + # Record ghost_col in schema_document["base_columns"] as if it was + # always a base column, but do NOT add it to _expected_base_fields + # (it is absent from the patched expected set below). + doc = cot.schema_document or {} + doc["base_columns"] = list(doc.get("base_columns", [])) + [ + {"name": "ghost_col", "field_class": "CharField", "null": True} + ] + CustomObjectType.objects.filter(pk=cot.pk).update(schema_document=doc) + cot.refresh_from_db() + + result = heal_cot(cot, verbosity=0) + + warned_types = [w["type"] for w in result["warned"]] + self.assertIn("removed_from_model", warned_types) + # Must not have tried to drop anything + self.assertEqual(result["added"], []) + + # Clean up + with connection.schema_editor() as editor: + editor.remove_field(cot.get_model(), ghost_field) + + def test_type_change_detected_as_warning(self): + """A column present in DB and model but with a changed field class must warn.""" + cot = self.create_custom_object_type(name="hc_type", slug="hc-type") + + # Seed schema_document to claim 'created' was originally an IntegerField + # (in reality it's a DateTimeField). heal_cot should detect the mismatch. + doc = cot.schema_document or {} + cols = {c["name"]: c for c in doc.get("base_columns", [])} + if "created" in cols: + cols["created"] = {"name": "created", "field_class": "IntegerField", "null": False} + else: + cols["created"] = {"name": "created", "field_class": "IntegerField", "null": False} + doc["base_columns"] = list(cols.values()) + CustomObjectType.objects.filter(pk=cot.pk).update(schema_document=doc) + cot.refresh_from_db() + + result = heal_cot(cot, verbosity=0) + + warned_types = [w["type"] for w in result["warned"]] + self.assertIn("type_changed", warned_types) + changed = next(w for w in result["warned"] if w["type"] == "type_changed") + self.assertEqual(changed["field"], "created") + + # ------------------------------------------------------------------ + # dry_run mode + # ------------------------------------------------------------------ + + def test_dry_run_does_not_modify_db(self): + """dry_run=True must report additions without touching the DB.""" + cot = self.create_custom_object_type(name="hc_dryrun", slug="hc-dryrun") + table_name = cot.get_database_table_name() + + from django.db import models as dj_models + extra_field = dj_models.CharField(max_length=10, null=True, blank=True) + extra_field.name = "dry_col" + extra_field.column = "dry_col" + extra_field.set_attributes_from_name("dry_col") + extra_field.model = cot.get_model() + + base_fields = _expected_base_fields(cot) + base_fields["dry_col"] = extra_field + + with patch( + "netbox_custom_objects.mixin_migration._expected_base_fields", + return_value=base_fields, + ): + result = heal_cot(cot, verbosity=0, dry_run=True) + + # Column must be reported as would-be-added + self.assertIn("dry_col", result["added"]) + + # But must NOT exist in the actual DB + with connection.cursor() as cur: + actual = { + c.name for c in connection.introspection.get_table_description(cur, table_name) + } + self.assertNotIn("dry_col", actual) + + def test_dry_run_does_not_update_snapshot(self): + """dry_run=True must not update schema_document.""" + cot = self.create_custom_object_type(name="hc_drysn", slug="hc-drysn") + original_doc = cot.schema_document + + from django.db import models as dj_models + extra_field = dj_models.CharField(max_length=10, null=True, blank=True) + extra_field.name = "dry_snap_col" + extra_field.column = "dry_snap_col" + extra_field.set_attributes_from_name("dry_snap_col") + extra_field.model = cot.get_model() + + base_fields = _expected_base_fields(cot) + base_fields["dry_snap_col"] = extra_field + + with patch( + "netbox_custom_objects.mixin_migration._expected_base_fields", + return_value=base_fields, + ): + heal_cot(cot, verbosity=0, dry_run=True) + + cot.refresh_from_db() + self.assertEqual(cot.schema_document, original_doc) + + +# --------------------------------------------------------------------------- +# heal_all_cots() +# --------------------------------------------------------------------------- + +class HealAllCotsTestCase( + TransactionCleanupMixin, CustomObjectsTestCase, TransactionTestCase +): + """Tests for heal_all_cots() summary behaviour.""" + + def test_summary_total_matches_cot_count(self): + for i in range(3): + self.create_custom_object_type(name=f"hac_{i}", slug=f"hac-{i}") + summary = heal_all_cots(verbosity=0) + self.assertGreaterEqual(summary["total"], 3) + + def test_summary_healed_zero_when_no_drift(self): + self.create_custom_object_type(name="hac_nd", slug="hac-nd") + summary = heal_all_cots(verbosity=0) + self.assertEqual(summary["healed"], 0) + self.assertEqual(summary["warnings"], 0) + + def test_summary_keys_present(self): + summary = heal_all_cots(verbosity=0) + self.assertIn("total", summary) + self.assertIn("healed", summary) + self.assertIn("warnings", summary) + + +# --------------------------------------------------------------------------- +# Management command +# --------------------------------------------------------------------------- + +class UpgradeCustomObjectsCommandTestCase( + TransactionCleanupMixin, CustomObjectsTestCase, TransactionTestCase +): + """Smoke tests for the upgrade_custom_objects management command.""" + + def _call_command(self, *args, **kwargs): + from django.core.management import call_command + from io import StringIO + out = StringIO() + err = StringIO() + call_command( + "upgrade_custom_objects", *args, stdout=out, stderr=err, **kwargs + ) + return out.getvalue(), err.getvalue() + + def test_command_runs_without_error(self): + self.create_custom_object_type(name="cmd_basic", slug="cmd-basic") + stdout, stderr = self._call_command(verbosity=0) + # No exception means success; no DB errors in stderr + self.assertNotIn("Error", stderr) + + def test_dry_run_flag_accepted(self): + self.create_custom_object_type(name="cmd_dry", slug="cmd-dry") + stdout, stderr = self._call_command("--dry-run", verbosity=1) + self.assertIn("DRY RUN", stdout) + + def test_cot_flag_by_name(self): + cot = self.create_custom_object_type(name="cmd_cot", slug="cmd-cot") + stdout, _ = self._call_command("--cot", cot.name, verbosity=1) + # Should succeed with no-drift message + self.assertIn("no drift detected", stdout) + + def test_cot_flag_by_id(self): + cot = self.create_custom_object_type(name="cmd_cotid", slug="cmd-cotid") + stdout, _ = self._call_command("--cot", str(cot.pk), verbosity=1) + self.assertIn("no drift detected", stdout) + + def test_unknown_cot_raises_error(self): + from django.core.management.base import CommandError + with self.assertRaises(CommandError): + self._call_command("--cot", "nonexistent_cot_xyz")