Merge pull request #33 from wherobots/peter/region

feat: add region to smoke tests

Author: peterfoldes

tests/smoke.py

@@ -13,11 +13,13 @@
 from wherobots.db import connect, connect_direct
 from wherobots.db.constants import DEFAULT_ENDPOINT
 from wherobots.db.connection import Connection
+from wherobots.db.region import Region
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--api-key-file", help="File containing the API key")
     parser.add_argument("--token-file", help="File containing the token")
+    parser.add_argument("--region", help="Region to connect to (ie. aws-us-west-2)")
     parser.add_argument(
         "--debug",
         help="Enable debug logging",
@@ -73,9 +75,10 @@
             api_key=api_key,
             shutdown_after_inactive_seconds=args.shutdown_after_inactive_seconds,
             wait_timeout=900,
+            region=Region(args.region) if args.region else Region.AWS_US_WEST_2,
         )
 
-    def render(results: pandas.DataFrame):
+    def render(results: pandas.DataFrame) -> None:
         table = Table()
         table.add_column("#")
         for column in results.columns:

wherobots/db/driver.py

@@ -56,10 +56,10 @@ def gen_user_agent_header():
 
 def connect(
     host: str = DEFAULT_ENDPOINT,
-    token: str = None,
-    api_key: str = None,
-    runtime: Runtime = None,
-    region: Region = None,
+    token: Union[str, None] = None,
+    api_key: Union[str, None] = None,
+    runtime: Union[Runtime, None] = None,
+    region: Union[Region, None] = None,
     wait_timeout: float = DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS,
     read_timeout: float = DEFAULT_READ_TIMEOUT_SECONDS,
     reuse_session: bool = DEFAULT_REUSE_SESSION,
@@ -163,7 +163,7 @@ def http_to_ws(uri: str) -> str:
 def connect_direct(
     uri: str,
     protocol: Version = PROTOCOL_VERSION,
-    headers: Dict[str, str] = None,
+    headers: Union[Dict[str, str], None] = None,
     read_timeout: float = DEFAULT_READ_TIMEOUT_SECONDS,
     results_format: Union[ResultsFormat, None] = None,
     data_compression: Union[DataCompression, None] = None,

wherobots/db/region.py

@@ -3,3 +3,4 @@
 
 class Region(Enum):
     AWS_US_WEST_2 = "aws-us-west-2"
+    AWS_EU_EAST_1 = "aws-eu-west-1"