From 57fd5eb57c37d58065145e0c3a953284f76a4190 Mon Sep 17 00:00:00 2001
From: richard-dds <richard-ctr@friends.dds.mil>
Date: Thu, 16 Aug 2018 15:30:04 -0400
Subject: [PATCH] Add regex validations for treasury_code and ba_code

---
 atst/forms/financial.py       |  9 ++++++---
 tests/forms/test_financial.py | 37 +++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/atst/forms/financial.py b/atst/forms/financial.py
index c923be97..21274e9b 100644
--- a/atst/forms/financial.py
+++ b/atst/forms/financial.py
@@ -1,7 +1,7 @@
 import re
 from wtforms.fields.html5 import EmailField
 from wtforms.fields import StringField, SelectField
-from wtforms.validators import Required, Email, InputRequired
+from wtforms.validators import Required, Email, InputRequired, Regexp
 
 from atst.domain.exceptions import NotFoundError
 from atst.domain.pe_numbers import PENumbers
@@ -21,6 +21,9 @@ PE_REGEX = re.compile(
     re.X,
 )
 
+TREASURY_CODE_REGEX = re.compile(r"^0*([1-9]{4}|[1-9]{6})$")
+
+BA_CODE_REGEX = re.compile(r"^0*[1-9]{2}\w?$")
 
 def suggest_pe_id(pe_id):
     suggestion = pe_id
@@ -77,9 +80,9 @@ class FinancialForm(ValidatedForm):
 
     pe_id = StringField("Program Element (PE) Number related to your request", validators=[Required()])
 
-    treasury_code = StringField("Program Treasury Code", validators=[Required()])
+    treasury_code = StringField("Program Treasury Code", validators=[Required(), Regexp(TREASURY_CODE_REGEX)])
 
-    ba_code = StringField("Program BA Code", validators=[Required()])
+    ba_code = StringField("Program BA Code", validators=[Required(), Regexp(BA_CODE_REGEX)])
 
     fname_co = StringField("Contracting Officer First Name", validators=[Required()])
     lname_co = StringField("Contracting Officer Last Name", validators=[Required()])
diff --git a/tests/forms/test_financial.py b/tests/forms/test_financial.py
index 62cea8ec..9bc82421 100644
--- a/tests/forms/test_financial.py
+++ b/tests/forms/test_financial.py
@@ -30,3 +30,40 @@ def test_funding_type_other_required_if_funding_type_is_other():
     form = FinancialForm(data=form_data)
     form.validate()
     assert "funding_type_other" in form.errors
+
+
+@pytest.mark.parametrize("input_,expected", [
+    ("1234", True),
+    ("123456", True),
+    ("0001234", True),
+    ("000123456", True),
+    ("12345", False),
+    ("00012345", False),
+    ("0001234567", False),
+    ("000000", False),
+])
+def test_treasury_code_validation(input_, expected):
+    form_data = {"treasury_code": input_}
+    form = FinancialForm(data=form_data)
+    form.validate()
+    is_valid = "treasury_code" not in form.errors
+
+    assert is_valid == expected
+
+
+@pytest.mark.parametrize("input_,expected", [
+    ("12", True),
+    ("00012", True),
+    ("12A", True),
+    ("000123", True),
+    ("00012A", True),
+    ("0001", False),
+    ("00012AB", False),
+])
+def test_ba_code_validation(input_, expected):
+    form_data = {"ba_code": input_}
+    form = FinancialForm(data=form_data)
+    form.validate()
+    is_valid = "ba_code" not in form.errors
+
+    assert is_valid == expected