tests: update test suites to cover new endpoints

Signed-off-by: mkelly <[email protected]>

tests/zia/test_locations.py

@@ -1,6 +1,6 @@
 import pytest
 import responses
-from box import Box
+from box import Box, BoxList
 from responses import matchers
 
 from tests.conftest import stub_sleep
@@ -66,6 +66,22 @@ def fixture_sub_locations():
     ]
 
 
+@pytest.fixture(name="geo_locations")
+def geo_location_data():
+    return {
+        "city_geo_id": 5375480,
+        "state_geo_id": 5332921,
+        "latitude": 37.3897,
+        "longitude": -122.0832,
+        "city_name": "Mountain View",
+        "state_name": "California",
+        "country_name": "United States",
+        "country_code": "US",
+        "postal_code": "94041",
+        "continent_code": "NA",
+    }
+
+
 @responses.activate
 @stub_sleep
 def test_list_locations_with_one_page(zia, paginated_items):
@@ -183,7 +199,7 @@ def test_get_location_by_id(zia, locations):
 def test_get_location_by_name_and_id(zia):
     # Passing location_id and location_name should result in a ValueError.
     with pytest.raises(ValueError):
-        resp = zia.locations.get_location(location_id="1", location_name="Test A")
+        zia.locations.get_location(location_id="1", location_name="Test A")
 
 
 @responses.activate
@@ -267,7 +283,7 @@ def test_update_location(zia, locations):
 
     resp = zia.locations.update_location("1", name="Updated Test")
 
-    assert isinstance(resp, dict)
+    assert isinstance(resp, Box)
     assert resp.id == 1
     assert resp.name == "Updated Test"
 
@@ -381,3 +397,55 @@ def test_list_sublocations(zia, sub_locations):
     assert isinstance(resp, list)
     assert len(resp) == 2
     assert resp[0].id == 1
+
+
+@responses.activate
+def test_get_geo_by_coordinates(zia, geo_locations):
+    responses.add(
+        responses.GET,
+        url="https://zsapi.zscaler.net/api/v1/region/byGeoCoordinates",
+        json=geo_locations,
+        status=200,
+    )
+    resp = zia.locations.get_geo_by_coordinates(37, -122)
+    assert isinstance(resp, Box)
+    assert resp.city_name == "Mountain View"
+
+
+@responses.activate
+def test_get_geo_by_ip(zia, geo_locations):
+    responses.add(
+        responses.GET,
+        url="https://zsapi.zscaler.net/api/v1/region/byIPAddress/8.8.8.8",
+        json=geo_locations,
+        status=200,
+    )
+    resp = zia.locations.get_geo_by_ip("8.8.8.8")
+    assert isinstance(resp, Box)
+    assert resp.city_name == "Mountain View"
+
+
+@stub_sleep
+@responses.activate
+def test_list_cities_by_name(zia):
+    list_cities_data = [
+        {"city": "San Jose", "state": "CA", "country": "US"},
+        {"city": "San Francisco", "state": "CA", "country": "US"},
+    ]
+
+    responses.add(
+        responses.GET,
+        url="https://zsapi.zscaler.net/api/v1/region/search?prefix=San&page=1",
+        json=list_cities_data,
+        status=200,
+    )
+    responses.add(
+        responses.GET,
+        url="https://zsapi.zscaler.net/api/v1/region/search?prefix=San&page=2",
+        json=[],
+        status=200,
+    )
+    resp = zia.locations.list_cities_by_name(prefix="San")
+    assert isinstance(resp, BoxList)
+    assert len(resp) == 2
+    assert resp[0].city == "San Jose"

tests/zia/test_traffic.py

@@ -215,6 +215,14 @@ def fixture_vpn_credentials():
     ]
 
 
+@pytest.fixture(name="ipv6_prefixes")
+def fixture_ipv64_prefixes():
+    return [
+        {"id": 0, "name": "sample1", "prefixMask": "mask1", "dnsPrefix": True},
+        {"id": 1, "name": "sample2", "prefixMask": "mask2", "dnsPrefix": False},
+    ]
+
+
 @responses.activate
 @stub_sleep
 def test_list_gre_tunnels(zia, gre_tunnels):
@@ -508,8 +516,8 @@ def test_get_vpn_credential_by_fqdn(zia, vpn_credentials):
 
 
 def test_get_vpn_credential_error(zia):
-    with pytest.raises(Exception) as e_info:
-        resp = zia.traffic.get_vpn_credential("1", "[email protected]")
+    with pytest.raises(Exception):
+        zia.traffic.get_vpn_credential("1", "[email protected]")
 
 
 @responses.activate
@@ -575,3 +583,54 @@ def test_list_vips(zia, vips):
     assert isinstance(resp, BoxList)
     assert len(resp) == 2
     assert resp[0].data_center == "TESTA"
+
+
+@responses.activate
+def test_list_dns64_prefixes(zia, ipv6_prefixes):
+    responses.add(
+        responses.GET,
+        url="https://zsapi.zscaler.net/api/v1/ipv6config/dns64prefix",
+        json=ipv6_prefixes,
+        status=200,
+    )
+    resp = zia.traffic.list_dns64_prefixes()
+    assert isinstance(resp, BoxList)
+    assert len(resp) == 2
+
+
+@stub_sleep
+@responses.activate
+def test_list_nat64_prefixes(zia, ipv6_prefixes):
+    responses.add(
+        responses.GET,
+        url="https://zsapi.zscaler.net/api/v1/ipv6config/nat64prefix?page=1",
+        json=ipv6_prefixes,
+        status=200,
+    )
+    responses.add(
+        responses.GET,
+        url="https://zsapi.zscaler.net/api/v1/ipv6config/nat64prefix?page=2",
+        json=[],
+        status=200,
+    )
+    resp = zia.traffic.list_nat64_prefixes()
+    assert isinstance(resp, BoxList)
+    assert len(resp) == 2
+
+
+@responses.activate
+def test_list_gre_ip_addresses(zia):
+    gre_ip_addresses_data = [
+        {"ipAddress": "192.168.1.1", "greEnabled": True, "greTunnelIP": "10.0.0.1"},
+        {"ipAddress": "192.168.1.2", "greEnabled": False, "greTunnelIP": "10.0.0.2"},
+    ]
+
+    responses.add(
+        responses.GET,
+        url="https://zsapi.zscaler.net/api/v1/orgProvisioning/ipGreTunnelInfo",
+        json=gre_ip_addresses_data,
+        status=200,
+    )
+    resp = zia.traffic.list_gre_ip_addresses()
+    assert isinstance(resp, BoxList)
+    assert len(resp) == 2