Add test for iterator pattern

patterns/chapter_09_iterator/iterators.py

@@ -2,16 +2,34 @@
 
 from __future__ import annotations
 
+from abc import ABC
+from typing import Any
 
-class Iterator:
-    """Iterator class."""
+
+class Iterator(ABC):
+    """Abstract iterator class."""
+
+    def __iter__(self) -> Iterator:
+        """Return iterator."""
+        return self
+
+    def __next__(self) -> Any:
+        """Return next element (Abstract class).
+
+        All child classes must implement this method.
+        """
+        ...
+
+
+class MenuIterator:
+    """Menu iterator class."""
 
     def __init__(self, collection: list[tuple[str, float]]) -> None:
         """Initialize iterator."""
         self._collection = collection
         self._index = 0
 
-    def __iter__(self) -> Iterator:
+    def __iter__(self) -> MenuIterator:
         """Return iterator."""
         return self
 

patterns/chapter_09_iterator/main.py

@@ -9,13 +9,10 @@
 def run_pattern_example() -> None:
     """Run pattern example."""
     pancake_menu: Menu = PancackeMenu()
-    diner_menu: Menu = DinerMenu()
-
-    pancake_menu.add_item({"pancake": 1.99})
-    pancake_menu.add_item({"waffle": 2.99})
+    pancake_menu.set_items({"pancake": 1.99, "waffle": 2.99})
 
-    diner_menu.add_item(("spaghetti", 3.99))
-    diner_menu.add_item(("meatballs", 4.99))
+    diner_menu: Menu = DinerMenu()
+    diner_menu.set_items([("spaghetti", 3.99), ("meatballs", 4.99)])
 
     waitress: Waitress = Waitress([pancake_menu, diner_menu])
     waitress.print_all_menus()

patterns/chapter_09_iterator/menus.py

@@ -2,7 +2,7 @@
 
 from abc import ABC, abstractmethod
 
-from .iterators import Iterator
+from .iterators import MenuIterator
 
 
 class Menu(ABC):
@@ -19,7 +19,15 @@ def add_item(self, item) -> None:
         ...
 
     @abstractmethod
-    def create_iterator(self) -> Iterator:
+    def set_items(self, items) -> None:
+        """Set all menu items (abstract method).
+
+        All child classes must implement this method.
+        """
+        ...
+
+    @abstractmethod
+    def create_iterator(self) -> MenuIterator:
         """Create iterator for menu items (abstract method).
 
         All child classes must implement this method.
@@ -31,6 +39,14 @@ def description(self) -> str:
         """Return menu description."""
         return self._description
 
+    @abstractmethod
+    def __repr__(self) -> str:
+        """Return menu representation.
+
+        All child classes must implement this method.
+        """
+        ...
+
 
 class DinerMenu(Menu):
     """Diner menu class.
@@ -48,9 +64,17 @@ def add_item(self, item: tuple[str, float]) -> None:
         """Add item to menu."""
         self._menu_items.append(item)
 
-    def create_iterator(self) -> Iterator:
+    def set_items(self, items: list[tuple[str, float]]) -> None:
+        """Set all menu items."""
+        self._menu_items = items
+
+    def create_iterator(self) -> MenuIterator:
         """Create iterator for menu items."""
-        return Iterator(self._menu_items)
+        return MenuIterator(self._menu_items)
+
+    def __repr__(self) -> str:
+        """Return menu representation."""
+        return f'{self._description} with {len(self._menu_items)} items'
 
 
 class PancackeMenu(Menu):
@@ -69,6 +93,14 @@ def add_item(self, item: dict[str, float]) -> None:
         """Add item to menu."""
         self._menu_items.update(item)
 
-    def create_iterator(self) -> Iterator:
+    def set_items(self, items: dict[str, float]) -> None:
+        """Set all menu items."""
+        self._menu_items = items
+
+    def create_iterator(self) -> MenuIterator:
         """Create iterator for menu items."""
-        return Iterator(list(self._menu_items.items()))
+        return MenuIterator(list(self._menu_items.items()))
+
+    def __repr__(self) -> str:
+        """Return menu representation."""
+        return f'{self._description} with {len(self._menu_items)} items'

patterns/chapter_09_iterator/waitress.py

@@ -2,7 +2,7 @@
 
 from typing import Iterable, cast
 
-from .iterators import Iterator
+from .iterators import MenuIterator
 from .menus import Menu
 
 
@@ -16,11 +16,11 @@ def __init__(self, menus: Iterable[Menu]) -> None:
     def print_all_menus(self) -> None:
         """Print all menus."""
         for menu in self._menus:
-            print(menu.description)
+            print(menu)
             iterator = menu.create_iterator()
             self._print_menu(iterator)
 
-    def _print_menu(self, iterator: Iterator) -> None:
+    def _print_menu(self, iterator: MenuIterator) -> None:
         """Print one menu."""
         for item, price, *_ in iterator:
             print(f'- {cast(str, item)}: {price}')

tests/test_chapter_09/test_iterator_pattern.py

@@ -0,0 +1,33 @@
+"""Module for testing pattern "Iterator"."""
+
+
+import pytest
+
+from patterns.chapter_09_iterator.menus import DinerMenu, Menu, PancackeMenu
+from patterns.chapter_09_iterator.waitress import Waitress
+
+pytestmark = pytest.mark.parametrize(
+    ('menu', 'menu_items'),
+    [
+        (PancackeMenu(), {'pancake': 1.99, 'waffle': 2.99}),
+        (DinerMenu(), [('spaghetti', 3.99)]),
+    ],
+)
+
+
+def test_iterator_pattern(
+    capsys,
+    menu: Menu,
+    menu_items: list[tuple[str, float]] | dict[str, float],
+) -> None:
+    """Test iterator pattern.
+
+    Printing all menus must be unified even if menu_items are of different
+    types.
+    """
+    menu.set_items(menu_items)
+
+    Waitress([menu]).print_all_menus()
+    captured = capsys.readouterr()
+
+    assert str(menu) in captured.out