From c135039eb161961d9852d87c09d7bda3e4d84db0 Mon Sep 17 00:00:00 2001 From: penguinboi Date: Thu, 30 Jul 2026 23:19:11 -0400 Subject: [PATCH] Add database guard for generators --- grace/database.py | 7 +++ grace/generators/migration_generator.py | 13 ++++-- grace/generators/model_generator.py | 15 +++++-- tests/generators/test_migration_generator.py | 43 ++++++++++++++++++ tests/generators/test_model_generator.py | 46 ++++++++++++++++++++ tests/test_database.py | 23 ++++++++++ 6 files changed, 140 insertions(+), 7 deletions(-) create mode 100644 tests/generators/test_migration_generator.py create mode 100644 tests/generators/test_model_generator.py create mode 100644 tests/test_database.py diff --git a/grace/database.py b/grace/database.py index f69b65c..e98aba8 100644 --- a/grace/database.py +++ b/grace/database.py @@ -5,9 +5,16 @@ from alembic.util.exc import CommandError from .application import Application +from .exceptions import ConfigError def generate_migration(app: Application, message: str): + if not app.has_database: + raise ConfigError( + "This project has no database configured. " + "Run 'grace generate database' to add one." + ) + try: alembic_cfg = Config("alembic.ini") alembic_cfg.config_ini_section = app.environment diff --git a/grace/generators/migration_generator.py b/grace/generators/migration_generator.py index 316b747..e385f69 100644 --- a/grace/generators/migration_generator.py +++ b/grace/generators/migration_generator.py @@ -1,4 +1,4 @@ -from logging import info +from logging import info, warning from click.core import Argument @@ -22,11 +22,18 @@ def generate(self, message: str): grace generate migration "Add Greeting model" ``` """ - info(f"Generating migration '{message}'") - if not self.app: raise ValueError("app is not initialized") + if not self.app.has_database: + warning( + "This project has no database configured. " + "Run 'grace generate database' to add one." + ) + return + + info(f"Generating migration '{message}'") + generate_migration(self.app, message) def validate(self, message: str, **_kwargs) -> bool: diff --git a/grace/generators/model_generator.py b/grace/generators/model_generator.py index c4543f6..39e7400 100644 --- a/grace/generators/model_generator.py +++ b/grace/generators/model_generator.py @@ -1,4 +1,4 @@ -from logging import info +from logging import info, warning from re import match from click.core import Argument @@ -33,6 +33,16 @@ def generate(self, name: str, params: tuple[str]): grace generate model Greeting message:String lang:String ``` """ + if not self.app: + raise ValueError("app is not initialized") + + if not self.app.has_database: + warning( + "This project has no database configured. " + "Run 'grace generate database' to add one." + ) + return + info(f"Generating model '{name}'") columns, types = self.extract_columns(params) @@ -49,9 +59,6 @@ def generate(self, name: str, params: tuple[str]): output_dir="bot/models", ) - if not self.app: - raise ValueError("app is not initialized") - generate_migration(self.app, f"Create {name}") def validate(self, name: str, **_kwargs) -> bool: diff --git a/tests/generators/test_migration_generator.py b/tests/generators/test_migration_generator.py new file mode 100644 index 0000000..4727ad1 --- /dev/null +++ b/tests/generators/test_migration_generator.py @@ -0,0 +1,43 @@ +from unittest.mock import MagicMock + +import pytest + +from grace.generators.migration_generator import MigrationGenerator + + +@pytest.fixture +def generator(): + generator = MigrationGenerator() + generator.app = MagicMock(has_database=True) + return generator + + +def test_generate_without_app__expect_value_error(generator): + generator.app = None + + with pytest.raises(ValueError): + generator.generate("Add Greeting model") + + +def test_generate_without_database__expect_warning_and_no_migration( + mocker, generator, caplog +): + generator.app.has_database = False + mock_generate_migration = mocker.patch( + "grace.generators.migration_generator.generate_migration" + ) + + generator.generate("Add Greeting model") + + mock_generate_migration.assert_not_called() + assert "no database configured" in caplog.text.lower() + + +def test_generate_with_database__expect_migration_generated(mocker, generator): + mock_generate_migration = mocker.patch( + "grace.generators.migration_generator.generate_migration" + ) + + generator.generate("Add Greeting model") + + mock_generate_migration.assert_called_once_with(generator.app, "Add Greeting model") diff --git a/tests/generators/test_model_generator.py b/tests/generators/test_model_generator.py new file mode 100644 index 0000000..fbd3588 --- /dev/null +++ b/tests/generators/test_model_generator.py @@ -0,0 +1,46 @@ +from unittest.mock import MagicMock + +import pytest + +from grace.generator import Generator +from grace.generators.model_generator import ModelGenerator + + +@pytest.fixture +def generator(): + generator = ModelGenerator() + generator.app = MagicMock(has_database=True) + return generator + + +def test_generate_without_app__expect_value_error(generator): + generator.app = None + + with pytest.raises(ValueError): + generator.generate("Greeting", ("message:String",)) + + +def test_generate_without_database__expect_no_file_and_warning( + mocker, generator, caplog +): + generator.app.has_database = False + mock_generate_file = mocker.patch.object(Generator, "generate_file") + + generator.generate("Greeting", ("message:String",)) + + mock_generate_file.assert_not_called() + assert "no database configured" in caplog.text.lower() + + +def test_generate_with_database__expect_model_file_and_migration_generated( + mocker, generator +): + mock_generate_file = mocker.patch.object(Generator, "generate_file") + mock_generate_migration = mocker.patch( + "grace.generators.model_generator.generate_migration" + ) + + generator.generate("Greeting", ("message:String",)) + + mock_generate_file.assert_called_once() + mock_generate_migration.assert_called_once_with(generator.app, "Create Greeting") diff --git a/tests/test_database.py b/tests/test_database.py new file mode 100644 index 0000000..305e068 --- /dev/null +++ b/tests/test_database.py @@ -0,0 +1,23 @@ +from unittest.mock import MagicMock + +import pytest + +from grace.database import generate_migration +from grace.exceptions import ConfigError + + +def test_generate_migration_without_database__expect_config_error(): + app = MagicMock(has_database=False) + + with pytest.raises(ConfigError): + generate_migration(app, "Add Greeting model") + + +def test_generate_migration_with_database__expect_revision_called(mocker): + app = MagicMock(has_database=True, environment="development") + mock_revision = mocker.patch("grace.database.revision") + mocker.patch("grace.database.Config") + + generate_migration(app, "Add Greeting model") + + mock_revision.assert_called_once()