From dfaafb3014cdc0b28f07fa4e0faebf32351a9a98 Mon Sep 17 00:00:00 2001
From: tripleee <tripleee@users.noreply.github.com>
Date: Thu, 12 Dec 2019 20:36:23 +0200
Subject: [PATCH] test/test_blacklists.py: simple YAML validation

---
 test/test_blacklists.py | 64 ++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 63 insertions(+), 1 deletion(-)

diff --git a/test/test_blacklists.py b/test/test_blacklists.py
index 4f3d5c9536..55c21487a4 100755
--- a/test/test_blacklists.py
+++ b/test/test_blacklists.py
@@ -1,9 +1,12 @@
 #!/usr/bin/env python3
 # coding=utf-8
 
+import yaml
+from os import unlink
+
 import pytest
 
-from glob import glob
+from blacklists import Blacklist, YAMLParserCIDR, YAMLParserASN
 from helpers import files_changed, blacklist_integrity_check
 
 
@@ -22,3 +25,62 @@ def test_remote_diff():
     false_diff = "h j q t"
     assert files_changed(true_diff, file_set)
     assert not files_changed(false_diff, file_set)
+
+
+def test_yaml_blacklist():
+    with open('test_ip.yml', 'w') as y:
+        yaml.dump({
+            'Schema': 'yaml_cidr',
+            'Schema_version': '2019120601',
+            'items': [
+                {'ip': '1.2.3.4'},
+                {'ip': '2.3.4.5', 'disable': True},
+                {'ip': '3.4.5.6', 'comment': 'comment'},
+            ]}, y)
+    blacklist = Blacklist(('test_ip.yml', YAMLParserCIDR))
+    with pytest.raises(ValueError) as e:
+        blacklist.add('1.3.34')
+    with pytest.raises(ValueError) as e:
+        blacklist.add({'ip': '1.3.4'})
+    with pytest.raises(ValueError) as e:
+        blacklist.add({'ip': '1.2.3.4'})
+    with pytest.raises(ValueError) as e:
+        blacklist.add({'ip': '2.3.4.5'})
+    with pytest.raises(ValueError) as e:
+        blacklist.remove({'ip': '34.45.56.67'})
+    blacklist.add({'ip': '1.3.4.5'})
+    assert '1.2.3.4' in blacklist.parse()
+    assert '2.3.4.5' not in blacklist.parse()
+    assert '3.4.5.6' in blacklist.parse()
+    blacklist.remove({'ip': '3.4.5.6'})
+    assert '3.4.5.6' not in blacklist.parse()
+    unlink('test_ip.yml')
+
+
+def test_yaml_asn():
+    with open('test_asn.yml', 'w') as y:
+        yaml.dump({
+            'Schema': 'yaml_asn',
+            'Schema_version': '2019120601',
+            'items': [
+                {'asn': '123'},
+                {'asn': '234', 'disable': True},
+                {'asn': '345', 'comment': 'comment'},
+            ]}, y)
+    blacklist = Blacklist(('test_asn.yml', YAMLParserASN))
+    with pytest.raises(ValueError) as e:
+        blacklist.add('123')
+    with pytest.raises(ValueError) as e:
+        blacklist.add({'asn': 'invalid'})
+    with pytest.raises(ValueError) as e:
+        blacklist.add({'asn': '123'})
+    with pytest.raises(ValueError) as e:
+        blacklist.add({'asn': '234'})
+    with pytest.raises(ValueError) as e:
+        blacklist.remove({'asn': '9897'})
+    assert '123' in blacklist.parse()
+    assert '234' not in blacklist.parse()
+    assert '345' in blacklist.parse()
+    blacklist.remove({'asn': '345'})
+    assert '345' not in blacklist.parse()
+    unlink('test_asn.yml')