First commit

tests/test_dig_aws.py

@@ -0,0 +1,77 @@
+import json
+import sys
+
+import digaws.digaws as digaws
+from digaws import __description__, __version__
+
+import pytest
+
+import tests
+
+
+@pytest.fixture
+def test_dig():
+    return digaws.DigAWS(json.loads(tests.AWS_IP_RANGES))
+
+
+def test_cli(capsys):
+    sys.argv = ['digaws', '-h']
+    try:
+        digaws.main()
+    except SystemExit as e:
+        out, _ = capsys.readouterr()
+        assert __description__ in out
+        assert e.code == 0
+
+
+def test_cli_version(capsys, mocker):
+    sys.argv = ['digaws', '--version']
+    try:
+        digaws.main()
+    except SystemExit as e:
+        out, _ = capsys.readouterr()
+        assert out == f'digaws {__version__}\n'
+        assert e.code == 0
+
+
+def test_cli_invocation(capsys, mocker):
+    sys.argv = ['digaws', '52.94.76.0/22', '2600:1f14:fff:f810:a1c1:f507:a2d1:2dd8',
+                '--output', 'json']
+    mocker.patch('digaws.digaws.get_aws_ip_ranges', return_value=json.loads(tests.AWS_IP_RANGES))
+    digaws.main()
+    out, _ = capsys.readouterr()
+
+    assert out == tests.RESPONSE_JSON_JOINED_PRINT
+
+
+def test_dig_aws_construct(test_dig):
+    assert test_dig.ip_prefixes == tests.AWS_IPV4_RANGES_OBJ
+    assert test_dig.ipv6_prefixes == tests.AWS_IPV6_RANGES_OBJ
+
+
+def test_lookup(test_dig):
+    assert str(test_dig._lookup_data('52.94.76.1')[0]['ip_prefix']) == '52.94.76.0/22'
+    assert str(test_dig._lookup_data('52.94.76.0/24')[0]['ip_prefix']) == '52.94.76.0/22'
+
+    input = '2600:1f14:fff:f810:a1c1:f507:a2d1:2dd8'
+    assert str(test_dig._lookup_data(input)[0]['ipv6_prefix']) == '2600:1f14:fff:f800::/53'
+    assert str(test_dig._lookup_data(input)[1]['ipv6_prefix']) == '2600:1f14::/35'
+    assert str(test_dig._lookup_data('2600:1f14::/36')[0]['ipv6_prefix']) == '2600:1f14::/35'
+
+    with pytest.raises(ValueError) as e:
+        test_dig.lookup('what are you talking about')
+        assert e.startswith('Wrong IP or CIDR format')
+
+
+def test_response_plain_print(test_dig, capsys):
+    test_dig.lookup('52.94.76.0/22').plain_print()
+    out, _ = capsys.readouterr()
+
+    assert out == tests.RESPONSE_PLAIN_PRINT
+
+
+def test_response_json_print(test_dig, capsys):
+    test_dig.lookup('2600:1f14:fff:f810:a1c1:f507:a2d1:2dd8').json_print()
+    out, _ = capsys.readouterr()
+
+    assert out == tests.RESPONSE_JSON_PRINT