Add argument "sample" to aks functions

.travis.yml

@@ -3,6 +3,7 @@ language: python
 python:
   - 3.5
   - 3.6
+  - 3.7
 
 install:
   - pip install -r requirements.txt -r requirements-dev.txt

CHANGELOG.md

@@ -1,8 +1,27 @@
 # Changelog
 
 ## [Unreleased][]
-[Unreleased]: https://github.com/chaostoolkit-incubator/chaostoolkit-azure/compare/0.4.0...HEAD
+[Unreleased]: https://github.com/chaostoolkit-incubator/chaostoolkit-azure/compare/0.5.0...HEAD
 
+- Added the burn_io feature: increase the I/O operations per seconds of the 
+hard drive for a time period (default time period is 1 minute). Works by 
+copying random data into a burn folder, which is deleted at the end of the 
+script.
+- Added the network_latency feature: disturb the network of the VM, adding some
+ latency for a time period (defaults to a 200 +/- 50ms latency for 1 minute).
+  Only works on Linux machines for now.
+- Supporting multiple Azure Cloud, such as AZURE_CHINA_CLOUD.
+- Code clean up and refactoring, moving up client initiation.
+
+
+## [0.5.0][] - 2019-07-05
+
+[0.5.0]: https://github.com/chaostoolkit-incubator/chaostoolkit-azure/compare/0.4.0...0.5.0
+
+### Added
+
+-  Added the FillDisk feature: Create a file of random data on the disk of the 
+VM for a time period (defaults to a 1GB file for 2 minutes).
 -  Fixed the new MS Azure REST API version 2019-04-01
 
 ## [0.4.0][] - 2019-04-15

README.md

@@ -70,9 +70,18 @@ There are two ways of doing this:
     ```json
     {
         "azure": {
-            "client_id": "AZURE_CLIENT_ID",
-            "client_secret": "AZURE_CLIENT_SECRET",
-            "tenant_id": "AZURE_TENANT_ID"
+            "client_id": {
+                "type": "env",
+                "key": "AZURE_CLIENT_ID"
+            },
+            "client_secret": {
+                "type": "env",
+                "key": "AZURE_CLIENT_SECRET"
+            },
+            "tenant_id": {
+                "type": "env",
+                "key": "AZURE_TENANT_ID"
+            }
         }
     }
     ```
@@ -88,6 +97,20 @@ There are two ways of doing this:
         }
     }
     ```
+
+    Also if you are not working with Public Global Azure, e.g. China Cloud
+    You can feed the cloud environment name as well.
+    Please refer to msrestazure.azure_cloud
+    ```json
+        {
+            "azure": {
+                "client_id": "xxxxxxx",
+                "client_secret": "*******",
+                "tenant_id": "@@@@@@@@@@@",
+                "azure_cloud": "AZURE_CHINA_CLOUD"
+            }
+        }
+    ```
 
     Additionally you need to provide the Azure subscription id.
 

chaosazure/__init__.py

@@ -2,15 +2,23 @@
 
 """Top-level package for chaostoolkit-azure."""
 import contextlib
+from typing import List
+
+from azure.mgmt.compute import ComputeManagementClient
+from azure.mgmt.resourcegraph import ResourceGraphClient
 from chaoslib.discovery import initialize_discovery_result, discover_actions, \
     discover_probes
-from chaoslib.types import Discovery, DiscoveredActivities, Secrets
+from chaoslib.types import Discovery, DiscoveredActivities, \
+    Secrets, Configuration
 from logzero import logger
+import msrestazure
 from msrestazure.azure_active_directory import ServicePrincipalCredentials
-from typing import List
+from msrestazure.azure_cloud import AZURE_PUBLIC_CLOUD, \
+    AZURE_US_GOV_CLOUD, AZURE_GERMAN_CLOUD, AZURE_CHINA_CLOUD
 
-__all__ = ["auth", "discover", "__version__"]
-__version__ = '0.4.0'
+__all__ = ["auth", "discover", "__version__", "init_client",
+           "init_resource_graph_client"]
+__version__ = '0.5.0'
 
 
 @contextlib.contextmanager
@@ -33,6 +41,19 @@ def auth(secrets: Secrets) -> ServicePrincipalCredentials:
         }
     }
     ```
+    Also if you are not working with Public Global Azure, e.g. China Cloud
+    You can feed the cloud environment name as well.
+    Please refer to msrestazure.azure_cloud
+    ```json
+    {
+        "azure": {
+            "client_id": "xxxxxxx",
+            "client_secret": "*******",
+            "tenant_id": "@@@@@@@@@@@",
+            "azure_cloud": "AZURE_CHINA_CLOUD"
+        }
+    }
+    ```
 
     The client_id, tenant_id, and client_secret content will be read
     from the specified local environment variables, e.g. `AZURE_CLIENT_ID`,
@@ -48,14 +69,27 @@ def auth(secrets: Secrets) -> ServicePrincipalCredentials:
         resource_client = ResourceManagementClient(cred, azure_subscription_id)
         compute_client = ComputeManagementClient(cred, azure_subscription_id)
     ```
+
+    Again, if you are not working with Public Azure Cloud,
+    and you set azure_cloud in secret,
+    this will pass one more parameter `base_url` to above function.
+    ```python
+        cloud = __get_cloud_env_by_name(creds['azure_cloud'])
+        client = ComputeManagementClient(
+            credentials=cred, subscription_id=subscription_id,
+                        base_url=cloud.endpoints.resource_manager)
+    ```
+
     """
     creds = dict(
-        azure_client_id=None, azure_client_secret=None, azure_tenant_id=None)
+        azure_client_id=None, azure_client_secret=None,
+        azure_tenant_id=None, azure_cloud=None)
 
     if secrets:
         creds["azure_client_id"] = secrets.get("client_id")
         creds["azure_client_secret"] = secrets.get("client_secret")
         creds["azure_tenant_id"] = secrets.get("tenant_id")
+        creds["azure_cloud"] = secrets.get("azure_cloud", "AZURE_PUBLIC_CLOUD")
 
     credentials = __get_credentials(creds)
 
@@ -74,14 +108,39 @@ def discover(discover_system: bool = True) -> Discovery:
     return discovery
 
 
+def init_client(secrets: Secrets, configuration: Configuration) \
+        -> ComputeManagementClient:
+    with auth(secrets) as cred:
+        subscription_id = configuration['azure']['subscription_id']
+        base_url = __get_cloud_env_by_name(
+            secrets.get("azure_cloud")).endpoints.resource_manager
+        client = ComputeManagementClient(
+            credentials=cred, subscription_id=subscription_id,
+            base_url=base_url)
+
+        return client
+
+
+def init_resource_graph_client(secrets: Secrets) -> ResourceGraphClient:
+    with auth(secrets) as cred:
+        base_url = __get_cloud_env_by_name(
+            secrets.get("azure_cloud")).endpoints.resource_manager
+        client = ResourceGraphClient(
+            credentials=cred,
+            base_url=base_url)
+
+        return client
+
+
 ###############################################################################
 # Private functions
 ###############################################################################
-def __get_credentials(creds):
+def __get_credentials(creds: dict) -> ServicePrincipalCredentials:
     credentials = ServicePrincipalCredentials(
         client_id=creds['azure_client_id'],
         secret=creds['azure_client_secret'],
-        tenant=creds['azure_tenant_id']
+        tenant=creds['azure_tenant_id'],
+        cloud_environment=__get_cloud_env_by_name(creds['azure_cloud'])
     )
     return credentials
 
@@ -98,3 +157,23 @@ def __load_exported_activities() -> List[DiscoveredActivities]:
     activities.extend(discover_actions("chaosazure.webapp.actions"))
     activities.extend(discover_probes("chaosazure.webapp.probes"))
     return activities
+
+
+def __get_cloud_env_by_name(cloud: str = None) \
+        -> msrestazure.azure_cloud.Cloud:
+    try:
+        if cloud is None:
+            return msrestazure.azure_cloud.AZURE_PUBLIC_CLOUD
+
+        cloud = cloud.strip()
+        if cloud == "AZURE_CHINA_CLOUD":
+            return msrestazure.azure_cloud.AZURE_CHINA_CLOUD
+        elif cloud == "AZURE_US_GOV_CLOUD":
+            return msrestazure.azure_cloud.AZURE_US_GOV_CLOUD
+        elif cloud == "AZURE_GERMAN_CLOUD":
+            return msrestazure.azure_cloud.AZURE_GERMAN_CLOUD
+        else:
+            return msrestazure.azure_cloud.AZURE_PUBLIC_CLOUD
+    except AttributeError:
+        logger.info("Invalid azure cloud name, use default AZURE_PUBLIC_CLOUD")
+        return msrestazure.azure_cloud.AZURE_PUBLIC_CLOUD

chaosazure/aks/actions.py

@@ -13,6 +13,7 @@
 
 
 def delete_node(filter: str = None,
+                sample: int = None,
                 configuration: Configuration = None,
                 secrets: Secrets = None):
     """
@@ -33,11 +34,12 @@ def delete_node(filter: str = None,
         "Start delete_node: configuration='{}', filter='{}'".format(
             configuration, filter))
 
-    query = node_resource_group_query(filter, configuration, secrets)
+    query = node_resource_group_query(filter, sample, configuration, secrets)
     delete_machines(query, configuration, secrets)
 
 
 def stop_node(filter: str = None,
+              sample: int = None,
               configuration: Configuration = None,
               secrets: Secrets = None):
     """
@@ -55,11 +57,12 @@ def stop_node(filter: str = None,
         "Start stop_node: configuration='{}', filter='{}'".format(
             configuration, filter))
 
-    query = node_resource_group_query(filter, configuration, secrets)
+    query = node_resource_group_query(filter, sample, configuration, secrets)
     stop_machines(query, configuration, secrets)
 
 
 def restart_node(filter: str = None,
+                 sample: int = None,
                  configuration: Configuration = None,
                  secrets: Secrets = None):
     """
@@ -77,14 +80,14 @@ def restart_node(filter: str = None,
         "Start restart_node: configuration='{}', filter='{}'".format(
             configuration, filter))
 
-    query = node_resource_group_query(filter, configuration, secrets)
+    query = node_resource_group_query(filter, sample, configuration, secrets)
     restart_machines(query, configuration, secrets)
 
 
 ###############################################################################
 # Private helper functions
 ###############################################################################
-def node_resource_group_query(query, configuration, secrets):
+def node_resource_group_query(query, sample, configuration, secrets):
     aks = fetch_resources(query, RES_TYPE_AKS, secrets, configuration)
     if not aks:
         logger.warning("No AKS clusters found")
@@ -95,4 +98,13 @@ def node_resource_group_query(query, configuration, secrets):
                 [x['name'] for x in aks]))
     choice = random.choice(aks)
     node_resource_group = choice['properties']['nodeResourceGroup']
-    return "where resourceGroup =~ '{}'".format(node_resource_group)
+
+    return format_query(sample, node_resource_group)
+
+
+def format_query(sample, node_resource_group):
+    if sample is None:
+        return "where resourceGroup =~ '{}'".format(node_resource_group)
+    else:
+        return "where resourceGroup =~ '{}' | sample {}".format(
+            node_resource_group, sample)