import json import sys import pytest import digaws.digaws as digaws import tests from digaws import __description__, __version__ @pytest.fixture def test_dig(): return digaws.DigAWS( ip_ranges=json.loads(tests.AWS_IP_RANGES), output_fields=digaws.OUTPUT_FIELDS ) 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_cli_output_plain_fields_invocation(capsys, mocker): sys.argv = [ "digaws", "52.94.76.0/22", "--output=plain", "--output-fields", "region", ] mocker.patch("digaws.digaws.get_aws_ip_ranges", return_value=json.loads(tests.AWS_IP_RANGES)) digaws.main() out, _ = capsys.readouterr() assert out == "Region: us-west-2\n\n" def test_cli_output_json_fields_invocation(capsys, mocker): sys.argv = [ "digaws", "2600:1f14:fff:f810:a1c1:f507:a2d1:2dd8", "--output=json", "--output-fields", "service", "network_border_group", ] 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_FIELDS_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