diff --git a/edgar_tool/url_generator.py b/edgar_tool/url_generator.py index 181c8bb..61d3bfe 100644 --- a/edgar_tool/url_generator.py +++ b/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): diff --git a/tests/test_url_generator.py b/tests/test_url_generator.py index ad16902..77e5f45 100644 --- a/tests/test_url_generator.py +++ b/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.