Add test for invalid peo_in

edgar_tool/url_generator.py

@@ -67,8 +67,19 @@ def __init__(self, **query_args: SearchQueryKwargs):
         self.date_range_select = date_range_select
         self.start_date = start_date
         self.end_date = end_date
-        self.inc_in = query_args.get("inc_in")
-        self.peo_in = query_args.get("peo_in")
+
+        peo_in = query_args.get("peo_in")
+        if peo_in and peo_in not in PEO_IN_AND_INC_IN_TO_SEC_FORM_ID:
+            raise ValueError(
+                (
+                    "Invalid location code. "
+                    "Please provide a valid 2-letter state abbreviation, "
+                    "3-letter country code, or 'XX' for unknown."
+                )
+            )
+        inc_in = query_args.get("inc_in")
+        self.inc_in = inc_in
+        self.peo_in = peo_in
 
     @property
     def keywords(self):

tests/test_url_generator.py

@@ -544,4 +544,18 @@ def test_should_correctly_generate_search_url_for_principal_executive_office_in(
     assert actual_url == expected_url
 
 
+def test_should_raise_exception_if_location_code_invalid():
+    # GIVEN
+    expected_error_msg = (
+        "Invalid location code. "
+        "Please provide a valid 2-letter state abbreviation, "
+        "3-letter country code, or 'XX' for unknown."
+    )
+    test_kwargs = {"keywords": ["a"], "peo_in": "SUN"}
+
+    # WHEN / THEN
+    with pytest.raises(ValueError, match=expected_error_msg):
+        url_generator.generate_search_url_for_kwargs(test_kwargs)
+
+
 # TODO: Test incorporated in parameters.