diff --git a/tests/smoke.py b/tests/smoke.py index 5425944..699b400 100644 --- a/tests/smoke.py +++ b/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: diff --git a/wherobots/db/driver.py b/wherobots/db/driver.py index f8a770d..384bb11 100644 --- a/wherobots/db/driver.py +++ b/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, diff --git a/wherobots/db/region.py b/wherobots/db/region.py index f49325a..3673588 100644 --- a/wherobots/db/region.py +++ b/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"