import ast
import unittest
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
PLAYBOOK_SCRIPT = ROOT / "scripts" / "playbook.py"


class PlaybookTypingImportTests(unittest.TestCase):
    def test_optional_annotation_names_are_imported(self):
        tree = ast.parse(PLAYBOOK_SCRIPT.read_text(encoding="utf-8"))

        imported_names: set[str] = set()
        referenced_names: set[str] = set()

        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    imported_names.add(alias.asname or alias.name.split(".")[0])
            elif isinstance(node, ast.ImportFrom):
                for alias in node.names:
                    imported_names.add(alias.asname or alias.name)
            elif isinstance(node, ast.Name):
                referenced_names.add(node.id)

        if "Optional" in referenced_names:
            self.assertIn("Optional", imported_names)


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