update 2022 publisher

compendium_v2/publishers/survey_publisher_2022.py

@@ -40,7 +40,8 @@ JOIN sections s ON q.section_id = s.id
 JOIN compendia c ON s.compendium_id = c.id
 WHERE
     a.question_id = {}
-AND c.year = 2022
+    AND c.year = 2022
+    AND a.value NOT IN ('""', '"NA"', '"N/A"')
 ORDER BY n.id, a.question_id, a.updated_at DESC
 """
 
@@ -53,7 +54,7 @@ class FundingSource(enum.Enum):
     GOV_PUBLIC_BODIES = 16409
 
 
-class Question(enum.Enum):
+class StaffQuestion(enum.Enum):
     """
     Answers are numbers expressed in FTEs (full time equivalents)
     """
@@ -63,6 +64,13 @@ class Question(enum.Enum):
     non_technical_fte = 16417
 
 
+class OrgQuestion(enum.Enum):
+    """
+    Answers are strings
+    """
+    parent_org_name = 16419
+
+
 def query_budget():
     with survey_db.session_scope() as survey:
         return survey.execute(text(BUDGET_QUERY))
@@ -75,7 +83,7 @@ def query_funding_sources():
             yield source, survey.execute(text(query))
 
 
-def query_question(question: Question):
+def query_question(question: enum.Enum):
     query = QUESTION_TEMPLATE_QUERY.format(question.value)
     with survey_db.session_scope() as survey:
         return survey.execute(text(query))
@@ -170,7 +178,7 @@ def transfer_staff_data():
         nren_dict = helpers.get_uppercase_nren_dict(session)
 
         data = {}
-        for question in Question:
+        for question in StaffQuestion:
             rows = query_question(question)
             for row in rows:
                 nren_name = row[0].upper()
@@ -185,10 +193,10 @@ def transfer_staff_data():
                     continue
 
                 # initialize on first use, so we don't add data for nrens with no answers
-                data.setdefault(nren_name, {question: 0 for question in Question})[question] = value
+                data.setdefault(nren_name, {question: 0 for question in StaffQuestion})[question] = value
 
         for nren_name, nren_info in data.items():
-            if sum([nren_info[question] for question in Question]) == 0:
+            if sum([nren_info[question] for question in StaffQuestion]) == 0:
                 logger.info(f'{nren_name} has no staff data. Deleting if exists.')
                 session.query(model.NrenStaff).filter(
                     model.NrenStaff.nren_id == nren_dict[nren_name].id,
@@ -196,8 +204,8 @@ def transfer_staff_data():
                 ).delete()
                 continue
 
-            employed = nren_info[Question.permanent_fte] + nren_info[Question.subcontracted_fte]
-            technical = nren_info[Question.technical_fte] + nren_info[Question.non_technical_fte]
+            employed = nren_info[StaffQuestion.permanent_fte] + nren_info[StaffQuestion.subcontracted_fte]
+            technical = nren_info[StaffQuestion.technical_fte] + nren_info[StaffQuestion.non_technical_fte]
             if not math.isclose(
                 employed,
                 technical,
@@ -210,20 +218,52 @@ def transfer_staff_data():
             staff_data = model.NrenStaff(
                 nren_id=nren_dict[nren_name].id,
                 year=2022,
-                permanent_fte=nren_info[Question.permanent_fte],
-                subcontracted_fte=nren_info[Question.subcontracted_fte],
-                technical_fte=nren_info[Question.technical_fte],
-                non_technical_fte=nren_info[Question.non_technical_fte],
+                permanent_fte=nren_info[StaffQuestion.permanent_fte],
+                subcontracted_fte=nren_info[StaffQuestion.subcontracted_fte],
+                technical_fte=nren_info[StaffQuestion.technical_fte],
+                non_technical_fte=nren_info[StaffQuestion.non_technical_fte],
             )
             session.merge(staff_data)
         session.commit()
 
 
+def transfer_nren_parent_org():
+
+    # clean up the data a bit by removing some strings
+    strings_to_replace = [
+        'We are affiliated to '
+    ]
+
+    with db.session_scope() as session:
+        nren_dict = helpers.get_uppercase_nren_dict(session)
+
+        rows = query_question(OrgQuestion.parent_org_name)
+        for row in rows:
+            nren_name = row[0].upper()
+            value = str(row[1]).replace('"', '')
+
+            for string in strings_to_replace:
+                value = value.replace(string, '')
+
+            if nren_name not in nren_dict:
+                logger.info(f'{nren_name} unknown. Skipping.')
+                continue
+
+            parent_org = model.ParentOrganization(
+                nren_id=nren_dict[nren_name].id,
+                year=2022,
+                organization=value,
+            )
+            session.merge(parent_org)
+        session.commit()
+
+
 def _cli(config):
     helpers.init_db(config)
     transfer_budget()
     transfer_funding_sources()
     transfer_staff_data()
+    transfer_nren_parent_org()
 
 
 @click.command()

test/test_survey_publisher_2022.py

 from compendium_v2 import db
 from compendium_v2.db import model
-from compendium_v2.publishers.survey_publisher_2022 import _cli, FundingSource, Question
+from compendium_v2.publishers.survey_publisher_2022 import _cli, FundingSource, StaffQuestion, OrgQuestion
 
 
 def test_publisher(client, mocker, dummy_config):
@@ -31,33 +31,38 @@ def test_publisher(client, mocker, dummy_config):
         ]
 
     def question_data(question):
-        if question == Question.non_technical_fte:
+        if question == StaffQuestion.non_technical_fte:
             return [
                 ('nren1', '10'),
                 ('nren2', '80'),
                 ('nren3', '30'),
             ]
 
-        if question == Question.technical_fte:
+        if question == StaffQuestion.technical_fte:
             return [
                 ('nren1', '50'),
                 ('nren2', '20'),
                 ('nren3', '30'),
             ]
 
-        if question == Question.permanent_fte:
+        if question == StaffQuestion.permanent_fte:
             return [
                 ('nren1', '40'),
                 ('nren2', 'abc'),
                 ('nren3', '30'),
             ]
 
-        if question == Question.subcontracted_fte:
+        if question == StaffQuestion.subcontracted_fte:
             return [
                 ('nren1', '0'),
                 ('nren2', '0'),
                 ('nren3', '0'),
             ]
+        if question == OrgQuestion.parent_org_name:
+            return [
+                ('nren1', 'Org1'),
+                ('nren3', 'Org3'),
+            ]
 
     mocker.patch('compendium_v2.publishers.survey_publisher_2022.query_budget',
                  get_rows_as_tuples)
@@ -124,3 +129,13 @@ def test_publisher(client, mocker, dummy_config):
         assert staff_data[2].technical_fte == 30
         assert staff_data[2].permanent_fte == 30
         assert staff_data[2].subcontracted_fte == 0
+
+        org_data = session.query(model.ParentOrganization).order_by(
+            model.ParentOrganization.nren_id.asc()).all()
+
+        assert len(org_data) == 2
+        assert org_data[0].nren.name.lower() == 'nren1'
+        assert org_data[0].organization == 'Org1'
+
+        assert org_data[1].nren.name.lower() == 'nren3'
+        assert org_data[1].organization == 'Org3'