From 17d4d65c84d816d1d5bd08ca7a25f792113ee496 Mon Sep 17 00:00:00 2001
From: Laurent Mazuel <laurent.mazuel@gmail.com>
Date: Tue, 19 Dec 2017 17:33:08 -0800
Subject: [PATCH] Improve launcher

---
 example.py       | 21 ++++++++++++++-------
 samples/tools.py |  6 +++++-
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/example.py b/example.py
index f10e6e9..4233fe0 100644
--- a/example.py
+++ b/example.py
@@ -13,14 +13,21 @@
 import samples.tools
 
 def run_all_samples():
-    for _, sample_name, _ in pkgutil.walk_packages(samples.__path__):
-        sample_module = importlib.import_module('samples.'+sample_name)
-        subkey_env_name = getattr(sample_module, "SUBSCRIPTION_KEY_ENV_NAME", None)
-        if not subkey_env_name:
+    for _, section_name_name, ispkg in pkgutil.walk_packages(samples.__path__):
+        if not ispkg:
             continue
-        print("Executing sample from ", sample_name)
-        sample_module = importlib.import_module('samples.'+sample_name)
-        samples.tools.execute_samples(sample_module.__dict__, subkey_env_name)
+        section_package_name = "samples."+section_name_name
+        section_package = importlib.import_module(section_package_name)
+        for _, sample_name, _ in pkgutil.iter_modules(section_package.__path__):
+            sample_module = importlib.import_module(section_package_name+"."+sample_name)
+            subkey_env_name = getattr(sample_module, "SUBSCRIPTION_KEY_ENV_NAME", None)
+            if not subkey_env_name:
+                continue
+            print("Executing sample from ", sample_name)
+            try:
+                samples.tools.execute_samples(sample_module.__dict__, subkey_env_name)
+            except samples.tools.SubscriptionKeyError as err:
+                print("{}\n".format(err))
 
 if __name__ == "__main__":
     run_all_samples()
diff --git a/samples/tools.py b/samples/tools.py
index f92b0bc..563cba1 100644
--- a/samples/tools.py
+++ b/samples/tools.py
@@ -11,6 +11,10 @@
 import types
 
 
+class SubscriptionKeyError(Exception):
+    pass
+
+
 def start_sample(func, subscription_key):
     """Start the function and show its doc on output.
     """
@@ -25,7 +29,7 @@ def execute_samples(module_globals, key_env_variable):
     try:
         subscription_key = sys.argv[1] if len(sys.argv) >= 2 else os.environ[key_env_variable]
     except KeyError:
-        sys.exit("You need to either set the {} env variable or give it as env variable.".format(key_env_variable))
+        raise SubscriptionKeyError("You need to either set the {} env variable.".format(key_env_variable))
 
     for func in list(module_globals.values()):
         if not isinstance(func, types.FunctionType):