tests/test_stages.py: Fix test_json_zip issue

openownership · Jul 26, 2023 · a6097eb · a6097eb
1 parent fcb92aa
commit a6097eb
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/tests/test_stages.py b/tests/test_stages.py
@@ -137,9 +137,13 @@ def test_json_zip(self, temp_dir, output_dir, source_dir):
         with zipfile.ZipFile(output_dir / 'json.zip') as test_zip:
             with test_zip.open(f'{self.source}.json') as output_file:
                 data = output_file.readlines()
-                print(data)
-                assert len(data) == 20
-                assert json.loads(data[0].strip())['interestedParty']['describedByPersonStatement'] == '14105856581894595060'
+                statements = [json.loads(d.strip()) for d in data]
+                assert len(statements) == 20
+                assert len([s for s in statements if s["statementType"] == "ownershipOrControlStatement"]) == 16
+                assert '14105856581894595060' in [s['interestedParty']['describedByPersonStatement'] for s in statements if
+                                                                        s["statementType"] == "ownershipOrControlStatement"]
+                assert len([s for s in statements if s["statementType"] == "personStatement"]) == 4
+                assert '11262152698769124205' in [s["statementID"] for s in statements if s["statementType"] == "personStatement"]
 
     def test_sqlite_zip(self, temp_dir, output_dir, source_dir):
         """Test creation of output sqlite.db.gz file"""