reorganize post-generate tasks

Author: Galarzaa90

docs/intro.md

@@ -34,9 +34,12 @@ This fetches all the relevant articles from TibiaWiki and stores them in the dat
 
 It accepts the following parameters:
 
-- `-s`/`--skip-images` Option to skip fetching and saving images.
-- `-db`/ `--db-name` The name of the generated database file. `tibiawiki.db` by default.
-- `-sd`/ `--skip-deprecated` Option to skip deprecated articles when parsing.
+- `-i`/`--skip-images` Option to skip fetching and saving images.
+- `-o`/ `--db-name` The name of the generated database file. `tibiawiki.db` by default.
+- `-d`/ `--skip-deprecated` Option to skip deprecated articles when parsing.
+- `-c`/ `--skip-category` Option to skip one or more categories (repeatable), using internal category keys such as `achievements`, `items`, `creatures`, `houses`, or `charms`.
+
+If skipping a category would break a hard dependency for another category, the dependent category is skipped automatically and a warning is shown.
 
 The generated database is saved in the current directory, as well as a folder called `images` with all the fetched images.
 

tests/test_generation.py

@@ -1,17 +1,18 @@
 import datetime
 import sqlite3
 import unittest
-from unittest.mock import patch
+from unittest.mock import Mock, patch
+
+from click.testing import CliRunner
 
 from tests import load_resource
+from tibiawikisql import __main__ as cli_module
+from tibiawikisql import generation as generation_module
 from tibiawikisql.api import Article
-from tibiawikisql.generation import (
-    WEAPON_PROFICIENCY_NAME_ARTICLE,
-    WEAPON_PROFICIENCY_TABLES_ARTICLE,
-    generate_item_proficiency_perks,
-    wiki_client,
-)
+from tibiawikisql.generation import WEAPON_PROFICIENCY_NAME_ARTICLE, WEAPON_PROFICIENCY_TABLES_ARTICLE
 from tibiawikisql.schema import ItemProficiencyPerkTable, ItemTable
+from tibiawikisql.tasks.item_proficiency_perks import generate_item_proficiency_perks
+from tibiawikisql.tasks.loot_statistics import generate_loot_statistics
 
 
 class TestGeneration(unittest.TestCase):
@@ -46,11 +47,18 @@ def test_generate_item_proficiency_perks(self):
         )
         data_store = {"items_map": {"amber axe": 1, "amber cudgel": 2}}
 
-        with (
-            patch.object(wiki_client, "get_article", side_effect=[mapping_article, tables_article]),
-            patch("tibiawikisql.generation.click.echo") as mock_echo,
-        ):
-            generate_item_proficiency_perks(self.conn, data_store)
+        wiki_client = Mock()
+        wiki_client.get_article.side_effect = [mapping_article, tables_article]
+        mock_echo = Mock()
+        generate_item_proficiency_perks(
+            self.conn,
+            data_store,
+            wiki_client=wiki_client,
+            mapping_article_title=WEAPON_PROFICIENCY_NAME_ARTICLE,
+            tables_article_title=WEAPON_PROFICIENCY_TABLES_ARTICLE,
+            timed=generation_module.timed,
+            echo=mock_echo,
+        )
 
         rows = self.conn.execute(
             "SELECT item_id, proficiency_level, skill_image, icon, effect "
@@ -91,11 +99,18 @@ def test_generate_item_proficiency_perks_warn_and_continue(self):
         )
         data_store = {"items_map": {"amber axe": 1, "amber cudgel": 2}}
 
-        with (
-            patch.object(wiki_client, "get_article", side_effect=[mapping_article, tables_article]),
-            patch("tibiawikisql.generation.click.echo") as mock_echo,
-        ):
-            generate_item_proficiency_perks(self.conn, data_store)
+        wiki_client = Mock()
+        wiki_client.get_article.side_effect = [mapping_article, tables_article]
+        mock_echo = Mock()
+        generate_item_proficiency_perks(
+            self.conn,
+            data_store,
+            wiki_client=wiki_client,
+            mapping_article_title=WEAPON_PROFICIENCY_NAME_ARTICLE,
+            tables_article_title=WEAPON_PROFICIENCY_TABLES_ARTICLE,
+            timed=generation_module.timed,
+            echo=mock_echo,
+        )
 
         rows = self.conn.execute(
             "SELECT item_id, proficiency_level, skill_image, icon, effect FROM item_proficiency_perk",
@@ -118,20 +133,31 @@ def get_article_side_effect(title: str) -> Article | None:
                 return mapping_article
             return None
 
-        with (
-            patch.object(wiki_client, "get_article", side_effect=get_article_side_effect) as mock_get_article,
-            patch("tibiawikisql.generation.click.echo") as mock_echo,
-        ):
-            generate_item_proficiency_perks(self.conn, data_store)
+        wiki_client = Mock()
+        wiki_client.get_article.side_effect = get_article_side_effect
+        mock_echo = Mock()
+        generate_item_proficiency_perks(
+            self.conn,
+            data_store,
+            wiki_client=wiki_client,
+            mapping_article_title=WEAPON_PROFICIENCY_NAME_ARTICLE,
+            tables_article_title=WEAPON_PROFICIENCY_TABLES_ARTICLE,
+            timed=generation_module.timed,
+            echo=mock_echo,
+        )
 
         rows = self.conn.execute("SELECT COUNT(*) FROM item_proficiency_perk").fetchone()
         self.assertEqual(0, rows[0])
         self.assertEqual(
-            [WEAPON_PROFICIENCY_NAME_ARTICLE, WEAPON_PROFICIENCY_TABLES_ARTICLE],
-            [call.args[0] for call in mock_get_article.call_args_list],
+            [
+                WEAPON_PROFICIENCY_NAME_ARTICLE,
+                WEAPON_PROFICIENCY_TABLES_ARTICLE,
+            ],
+            [call.args[0] for call in wiki_client.get_article.call_args_list],
         )
         messages = " ".join(call.args[0] for call in mock_echo.call_args_list if call.args)
         self.assertIn("Could not fetch weapon proficiency pages", messages)
+        self.assertIn(WEAPON_PROFICIENCY_TABLES_ARTICLE, messages)
 
     def test_generate_item_proficiency_perks_case_insensitive_section_match(self):
         timestamp = datetime.datetime.fromisoformat("2024-01-01T00:00:00+00:00")
@@ -154,13 +180,122 @@ def test_generate_item_proficiency_perks_case_insensitive_section_match(self):
         )
         data_store = {"items_map": {"stale bread of ancientness": 3}}
 
-        with (
-            patch.object(wiki_client, "get_article", side_effect=[mapping_article, tables_article]),
-            patch("tibiawikisql.generation.click.echo"),
-        ):
-            generate_item_proficiency_perks(self.conn, data_store)
+        wiki_client = Mock()
+        wiki_client.get_article.side_effect = [mapping_article, tables_article]
+        generate_item_proficiency_perks(
+            self.conn,
+            data_store,
+            wiki_client=wiki_client,
+            mapping_article_title=WEAPON_PROFICIENCY_NAME_ARTICLE,
+            tables_article_title=WEAPON_PROFICIENCY_TABLES_ARTICLE,
+            timed=generation_module.timed,
+            echo=Mock(),
+        )
 
         rows = self.conn.execute(
             "SELECT item_id, proficiency_level, skill_image, icon, effect FROM item_proficiency_perk",
         ).fetchall()
         self.assertEqual([(3, 1, "Club Skill Bonus", None, "+1 Club Fighting")], [tuple(row) for row in rows])
+
+
+class TestGenerationOrchestration(unittest.TestCase):
+    def setUp(self):
+        self.conn = sqlite3.connect(":memory:")
+
+    def tearDown(self):
+        self.conn.close()
+
+    def test_skip_category_excludes_fetch_and_parse(self):
+        with (
+            patch("tibiawikisql.generation.fetch_category_entries", return_value=[]) as mock_fetch,
+            patch.object(generation_module.wiki_client, "get_articles", return_value=[]) as mock_get_articles,
+            patch("tibiawikisql.generation.POST_TASKS", ()),
+        ):
+            generation_module.generate(self.conn, skip_categories=("achievements",))
+
+        fetched_categories = [call.args[0] for call in mock_fetch.call_args_list]
+        self.assertNotIn("Achievements", fetched_categories)
+        self.assertEqual(len(generation_module.CATEGORIES) - 1, len(mock_get_articles.call_args_list))
+
+    def test_hard_dependencies_auto_skip_categories(self):
+        with (
+            patch("tibiawikisql.generation.fetch_category_entries", return_value=[]) as mock_fetch,
+            patch.object(generation_module.wiki_client, "get_articles", return_value=[]),
+            patch("tibiawikisql.generation.POST_TASKS", ()),
+            patch("tibiawikisql.generation.click.echo") as mock_echo,
+        ):
+            generation_module.generate(self.conn, skip_categories=("items",))
+
+        fetched_categories = [call.args[0] for call in mock_fetch.call_args_list]
+        self.assertNotIn("Objects", fetched_categories)
+        self.assertNotIn("Keys", fetched_categories)
+        messages = " ".join(call.args[0] for call in mock_echo.call_args_list if call.args)
+        self.assertIn("Skipping category 'keys'", messages)
+
+    def test_post_tasks_are_gated_by_enabled_categories(self):
+        gated_task = Mock()
+        always_task = Mock()
+        post_tasks = (
+            generation_module.PostTask("needs_items", lambda *_: gated_task(), dependencies=("items",)),
+            generation_module.PostTask("always_runs", lambda *_: always_task()),
+        )
+        with (
+            patch("tibiawikisql.generation.fetch_category_entries", return_value=[]),
+            patch.object(generation_module.wiki_client, "get_articles", return_value=[]),
+            patch("tibiawikisql.generation.POST_TASKS", post_tasks),
+            patch("tibiawikisql.generation.click.echo") as mock_echo,
+        ):
+            generation_module.generate(self.conn, skip_categories=("items",))
+
+        gated_task.assert_not_called()
+        always_task.assert_called_once_with()
+        messages = " ".join(call.args[0] for call in mock_echo.call_args_list if call.args)
+        self.assertIn("Skipping task 'needs_items'", messages)
+
+    def test_generate_loot_statistics_early_return_without_maps(self):
+        wiki_client = Mock()
+        generate_loot_statistics(
+            self.conn,
+            {},
+            wiki_client=wiki_client,
+            progress_bar=generation_module.progress_bar,
+            article_label=generation_module.article_label,
+            timed=generation_module.timed,
+            echo=Mock(),
+        )
+        wiki_client.get_articles.assert_not_called()
+
+
+class TestGenerateCommand(unittest.TestCase):
+    def setUp(self):
+        self.runner = CliRunner()
+
+    def test_skip_category_option_is_passed_to_generate(self):
+        with patch("tibiawikisql.__main__.generation.generate") as mock_generate:
+            result = self.runner.invoke(
+                cli_module.cli,
+                [
+                    "generate",
+                    "--db-name",
+                    ":memory:",
+                    "--skip-category",
+                    "achievements",
+                    "--skip-category",
+                    "items",
+                ],
+            )
+        self.assertEqual(0, result.exit_code, result.output)
+        self.assertEqual(("achievements", "items"), mock_generate.call_args.kwargs["skip_categories"])
+
+    def test_skip_category_rejects_invalid_value(self):
+        result = self.runner.invoke(
+            cli_module.cli,
+            [
+                "generate",
+                "--skip-category",
+                "invalid-category",
+            ],
+        )
+        self.assertNotEqual(0, result.exit_code)
+        self.assertIn("Invalid value", result.output)
+        self.assertIn("--skip-category", result.output)

tibiawikisql/__main__.py

@@ -19,14 +19,29 @@ def cli():
 
 
 @cli.command(name="generate")
-@click.option('-s', '--skip-images', help="Skip fetching and loading images to the database.", is_flag=True)
-@click.option('-db', '--db-name', help="Name for the database file.", default=DATABASE_FILE)
-@click.option('-sd', '--skip-deprecated', help="Skips fetching deprecated articles and their images.", is_flag=True)
-def generate(skip_images, db_name, skip_deprecated):
+@click.option('-i', '--skip-images', help="Skip fetching and loading images to the database.", is_flag=True)
+@click.option('-o', '--db-name', help="Name for the database file.", default=DATABASE_FILE)
+@click.option('-d', '--skip-deprecated', help="Skips fetching deprecated articles and their images.", is_flag=True)
+@click.option(
+    "-c",
+    "--skip-category",
+    "skip_categories",
+    multiple=True,
+    type=click.Choice(sorted(generation.CATEGORIES), case_sensitive=False),
+    help=(
+        "Skip specific categories. Can be repeated."
+    ),
+)
+def generate(skip_images, db_name, skip_deprecated, skip_categories):
     """Generates a database file."""
     with timed() as t:
         with sqlite3.connect(db_name) as conn:
-            generation.generate(conn, skip_images, skip_deprecated)
+            generation.generate(
+                conn,
+                skip_images=skip_images,
+                skip_deprecated=skip_deprecated,
+                skip_categories=skip_categories,
+            )
     click.echo(f"Command finished in {t.elapsed:.2f} seconds.")