From f518c9835c1623ada38c7a0683b056969cead796 Mon Sep 17 00:00:00 2001
From: Darek Lubomski <darek.lubomski@gmail.com>
Date: Tue, 17 Sep 2024 15:25:34 +0200
Subject: [PATCH] digits also possible

---
 iata_code_fetcher/fetcher.py |  4 ++--
 tests/test_fetcher.py        | 42 ------------------------------------
 2 files changed, 2 insertions(+), 44 deletions(-)

diff --git a/iata_code_fetcher/fetcher.py b/iata_code_fetcher/fetcher.py
index 4cd4949..c783b5c 100644
--- a/iata_code_fetcher/fetcher.py
+++ b/iata_code_fetcher/fetcher.py
@@ -4,7 +4,7 @@
 """
 
 import json
-from string import ascii_uppercase
+from string import ascii_uppercase, digits
 from itertools import product
 from typing import Generator, List, Dict
 from enum import Enum
@@ -48,7 +48,7 @@ def generate_codes(length: int) -> Generator[str, None, None]:
     :param length: Length of the IATA code (2 for carrier, 3 for airport).
     :return: Generator of code strings.
     """
-    return ("".join(letters) for letters in product(ascii_uppercase, repeat=length))
+    return ("".join(letters) for letters in product(ascii_uppercase + digits, repeat=length))
 
 
 def fetch_and_process_data(code: str, code_type: CodeType) -> List[Dict[str, str]]:
diff --git a/tests/test_fetcher.py b/tests/test_fetcher.py
index aad42c7..f330ba9 100644
--- a/tests/test_fetcher.py
+++ b/tests/test_fetcher.py
@@ -166,47 +166,5 @@ def test_process_and_save_data_airport(mock_get, mock_file, airport_response_moc
     )
 
 
-def test_generate_codes_for_two_letter_codes():
-    """
-    Test to ensure the function generates two-letter codes correctly
-    """
-    length = 2  # Test with two-letter codes
-    codes = list(generate_codes(length))
-    expected_number_of_codes = 26**2  # There are 26 letters, so 26^2 two-letter combinations
-
-    # Check if all codes have the correct length
-    assert all(len(code) == length for code in codes), "All codes must have the specified length of 2"
-
-    # Check the total number of generated codes
-    assert (
-        len(codes) == expected_number_of_codes
-    ), f"The number of generated two-letter codes should be {expected_number_of_codes}"
-
-    # Check specific codes to ensure correct sequence
-    assert codes[0] == "AA", "The first code should be 'AA'"
-    assert codes[-1] == "ZZ", "The last code should be 'ZZ'"
-
-
-def test_generate_codes_for_three_letter_codes():
-    """
-    Test to ensure the function generates three-letter codes correctly
-    """
-    length = 3  # Test with three-letter codes
-    codes = list(generate_codes(length))
-    expected_number_of_codes = 26**3  # There are 26 letters, so 26^3 three-letter combinations
-
-    # Check if all codes have the correct length
-    assert all(len(code) == length for code in codes), "All codes must have the specified length of 3"
-
-    # Check the total number of generated codes
-    assert (
-        len(codes) == expected_number_of_codes
-    ), f"The number of generated three-letter codes should be {expected_number_of_codes}"
-
-    # Check specific codes to ensure correct sequence
-    assert codes[0] == "AAA", "The first code should be 'AAA'"
-    assert codes[-1] == "ZZZ", "The last code should be 'ZZZ'"
-
-
 if __name__ == "__main__":
     pytest.main()