diff --git a/moto/__init__.py b/moto/__init__.py index 2f10b25dc..0481e49ee 100644 --- a/moto/__init__.py +++ b/moto/__init__.py @@ -127,6 +127,9 @@ mock_sts = lazy_load(".sts", "mock_sts") mock_sts_deprecated = lazy_load(".sts", "mock_sts_deprecated") mock_swf = lazy_load(".swf", "mock_swf") mock_swf_deprecated = lazy_load(".swf", "mock_swf_deprecated") +mock_timestreamwrite = lazy_load( + ".timestreamwrite", "mock_timestreamwrite", boto3_name="timestream-write" +) mock_transcribe = lazy_load(".transcribe", "mock_transcribe") XRaySegment = lazy_load(".xray", "XRaySegment") mock_xray = lazy_load(".xray", "mock_xray") diff --git a/moto/backend_index.py b/moto/backend_index.py index 8457a5c9e..a0401d42e 100644 --- a/moto/backend_index.py +++ b/moto/backend_index.py @@ -71,7 +71,7 @@ backend_url_patterns = [ re.compile("https?://.*\\.kinesisvideo\\.(.+)\\.amazonaws.com"), ), ("kms", re.compile("https?://kms\\.(.+)\\.amazonaws\\.com")), - ("lambda", re.compile("https?://lambda\\.(.+)\\.amazonaws\\.com(|.cn)")), + ("lambda", re.compile("https?://lambda\\.(.+)\\.amazonaws\\.com")), ("logs", re.compile("https?://logs\\.(.+)\\.amazonaws\\.com")), ( "managedblockchain", @@ -111,9 +111,17 @@ backend_url_patterns = [ ("ssm", re.compile("https?://ssm\\.(.+)\\.amazonaws\\.com")), ("ssm", re.compile("https?://ssm\\.(.+)\\.amazonaws\\.com\\.cn")), ("stepfunctions", re.compile("https?://states\\.(.+)\\.amazonaws.com")), - ("sts", re.compile("https?://sts\\.(.*\\.)?amazonaws\\.com(|.cn)")), + ("sts", re.compile("https?://sts\\.(.*\\.)?amazonaws\\.com")), ("support", re.compile("https?://support\\.(.+)\\.amazonaws\\.com")), ("swf", re.compile("https?://swf\\.(.+)\\.amazonaws\\.com")), + ( + "timestream-write", + re.compile("https?://ingest\\.timestream\\.(.+)\\.amazonaws\\.com"), + ), + ( + "timestream-write", + re.compile("https?://ingest\\.timestream\\.(.+)\\.amazonaws\\.com/"), + ), ("transcribe", re.compile("https?://transcribe\\.(.+)\\.amazonaws\\.com")), ("wafv2", re.compile("https?://wafv2\\.(.+)\\.amazonaws.com")), ("xray", re.compile("https?://xray\\.(.+)\\.amazonaws.com")), diff --git a/moto/server.py b/moto/server.py index 616d3bad9..dd293c957 100644 --- a/moto/server.py +++ b/moto/server.py @@ -56,6 +56,7 @@ class DomainDispatcherApplication(object): self.backend_url_patterns = backend_index.backend_url_patterns def get_backend_for_host(self, host): + if host == "moto_api": return host @@ -128,6 +129,10 @@ class DomainDispatcherApplication(object): host = "api.{service}.{region}.amazonaws.com".format( service=service, region=region ) + elif service == "timestream": + host = "ingest.{service}.{region}.amazonaws.com".format( + service=service, region=region + ) else: host = "{service}.{region}.amazonaws.com".format( service=service, region=region @@ -247,13 +252,17 @@ def create_backend_app(service): endpoint = original_endpoint + str(index) index += 1 - backend_app.add_url_rule( - url_path, - endpoint=endpoint, - methods=HTTP_METHODS, - view_func=view_func, - strict_slashes=False, - ) + # Some services do not provide a URL path + # I.e., boto3 sends a request to 'https://ingest.timestream.amazonaws.com' + # Which means we have a empty url_path to catch this request - but Flask can't handle that + if url_path: + backend_app.add_url_rule( + url_path, + endpoint=endpoint, + methods=HTTP_METHODS, + view_func=view_func, + strict_slashes=False, + ) backend_app.test_client_class = AWSTestHelper return backend_app diff --git a/moto/timestreamwrite/__init__.py b/moto/timestreamwrite/__init__.py new file mode 100644 index 000000000..f8b12f3bf --- /dev/null +++ b/moto/timestreamwrite/__init__.py @@ -0,0 +1,4 @@ +from .models import timestreamwrite_backends +from ..core.models import base_decorator + +mock_timestreamwrite = base_decorator(timestreamwrite_backends) diff --git a/moto/timestreamwrite/exceptions.py b/moto/timestreamwrite/exceptions.py new file mode 100644 index 000000000..3017dd382 --- /dev/null +++ b/moto/timestreamwrite/exceptions.py @@ -0,0 +1 @@ +"""Exceptions raised by the timestreamwrite service.""" diff --git a/moto/timestreamwrite/models.py b/moto/timestreamwrite/models.py new file mode 100644 index 000000000..614aa5d94 --- /dev/null +++ b/moto/timestreamwrite/models.py @@ -0,0 +1,180 @@ +from boto3 import Session + +from moto.core import ACCOUNT_ID, BaseBackend, BaseModel + + +class TimestreamTable(BaseModel): + def __init__(self, region_name, table_name, db_name, retention_properties): + self.region_name = region_name + self.name = table_name + self.db_name = db_name + self.retention_properties = retention_properties + self.records = [] + + def update(self, retention_properties): + self.retention_properties = retention_properties + + def write_records(self, records): + self.records.append(records) + + @property + def arn(self): + return f"arn:aws:timestream:{self.region_name}:{ACCOUNT_ID}:database/{self.db_name}/table/{self.name}" + + def description(self): + return { + "Arn": self.arn, + "TableName": self.name, + "DatabaseName": self.db_name, + "TableStatus": "ACTIVE", + "RetentionProperties": self.retention_properties, + } + + +class TimestreamDatabase(BaseModel): + def __init__(self, region_name, database_name, kms_key_id): + self.region_name = region_name + self.name = database_name + self.kms_key_id = kms_key_id + self.tables = dict() + + def update(self, kms_key_id): + self.kms_key_id = kms_key_id + + def create_table(self, table_name, retention_properties): + table = TimestreamTable( + region_name=self.region_name, + table_name=table_name, + db_name=self.name, + retention_properties=retention_properties, + ) + self.tables[table_name] = table + return table + + def update_table(self, table_name, retention_properties): + table = self.tables[table_name] + table.update(retention_properties=retention_properties) + return table + + def delete_table(self, table_name): + del self.tables[table_name] + + def describe_table(self, table_name): + return self.tables[table_name] + + def list_tables(self): + return self.tables.values() + + @property + def arn(self): + return ( + f"arn:aws:timestream:{self.region_name}:{ACCOUNT_ID}:database/{self.name}" + ) + + def description(self): + return { + "Arn": self.arn, + "DatabaseName": self.name, + "TableCount": len(self.tables.keys()), + "KmsKeyId": self.kms_key_id, + } + + +class TimestreamWriteBackend(BaseBackend): + def __init__(self, region_name): + self.region_name = region_name + self.databases = dict() + + def create_database(self, database_name, kms_key_id, tags): + database = TimestreamDatabase(self.region_name, database_name, kms_key_id) + self.databases[database_name] = database + return database + + def delete_database(self, database_name): + del self.databases[database_name] + + def describe_database(self, database_name): + return self.databases[database_name] + + def list_databases(self): + return self.databases.values() + + def update_database(self, database_name, kms_key_id): + database = self.databases[database_name] + database.update(kms_key_id=kms_key_id) + return database + + def create_table(self, database_name, table_name, retention_properties): + database = self.describe_database(database_name) + table = database.create_table(table_name, retention_properties) + return table + + def delete_table(self, database_name, table_name): + database = self.describe_database(database_name) + database.delete_table(table_name) + + def describe_table(self, database_name, table_name): + database = self.describe_database(database_name) + table = database.describe_table(table_name) + return table + + def list_tables(self, database_name): + database = self.describe_database(database_name) + tables = database.list_tables() + return tables + + def update_table(self, database_name, table_name, retention_properties): + database = self.describe_database(database_name) + table = database.update_table(table_name, retention_properties) + return table + + def write_records(self, database_name, table_name, records): + database = self.describe_database(database_name) + table = database.describe_table(table_name) + table.write_records(records) + + def describe_endpoints(self): + # https://docs.aws.amazon.com/timestream/latest/developerguide/Using-API.endpoint-discovery.how-it-works.html + # Usually, the address look like this: + # ingest-cell1.timestream.us-east-1.amazonaws.com + # Where 'cell1' can be any number, 'cell2', 'cell3', etc - whichever endpoint happens to be available for that particular account + # We don't implement a cellular architecture in Moto though, so let's keep it simple + return { + "Endpoints": [ + { + "Address": f"ingest.timestream.{self.region_name}.amazonaws.com", + "CachePeriodInMinutes": 1440, + } + ] + } + + def reset(self): + region_name = self.region_name + self.__dict__ = {} + self.__init__(region_name) + + +timestreamwrite_backends = {} +for available_region in Session().get_available_regions("timestream-write"): + timestreamwrite_backends[available_region] = TimestreamWriteBackend( + available_region + ) +for available_region in Session().get_available_regions( + "timestream-write", partition_name="aws-us-gov" +): + timestreamwrite_backends[available_region] = TimestreamWriteBackend( + available_region + ) +for available_region in Session().get_available_regions( + "timestream-write", partition_name="aws-cn" +): + timestreamwrite_backends[available_region] = TimestreamWriteBackend( + available_region + ) + +if len(timestreamwrite_backends) == 0: + # Boto does not return any regions at the time of writing (20/10/2021) + # Hardcoding the known regions for now + # Thanks, Jeff + for r in ["us-east-1", "us-east-2", "us-west-2", "eu-central-1", "eu-west-1"]: + timestreamwrite_backends[r] = TimestreamWriteBackend(r) diff --git a/moto/timestreamwrite/responses.py b/moto/timestreamwrite/responses.py new file mode 100644 index 000000000..84653d758 --- /dev/null +++ b/moto/timestreamwrite/responses.py @@ -0,0 +1,93 @@ +import json + +from moto.core.responses import BaseResponse +from .models import timestreamwrite_backends + + +class TimestreamWriteResponse(BaseResponse): + def __init__(self): + super().__init__() + + @property + def timestreamwrite_backend(self): + """Return backend instance specific for this region.""" + return timestreamwrite_backends[self.region] + + def create_database(self): + database_name = self._get_param("DatabaseName") + kms_key_id = self._get_param("KmsKeyId") + tags = self._get_list_prefix("Tags.member") + database = self.timestreamwrite_backend.create_database( + database_name=database_name, kms_key_id=kms_key_id, tags=tags, + ) + return json.dumps(dict(Database=database.description())) + + def delete_database(self): + database_name = self._get_param("DatabaseName") + self.timestreamwrite_backend.delete_database(database_name=database_name) + return "{}" + + def describe_database(self): + database_name = self._get_param("DatabaseName") + database = self.timestreamwrite_backend.describe_database( + database_name=database_name + ) + return json.dumps(dict(Database=database.description())) + + def update_database(self): + database_name = self._get_param("DatabaseName") + kms_key_id = self._get_param("KmsKeyId") + database = self.timestreamwrite_backend.update_database( + database_name, kms_key_id + ) + return json.dumps(dict(Database=database.description())) + + def list_databases(self): + all_dbs = self.timestreamwrite_backend.list_databases() + return json.dumps(dict(Databases=[db.description() for db in all_dbs])) + + def create_table(self): + database_name = self._get_param("DatabaseName") + table_name = self._get_param("TableName") + retention_properties = self._get_param("RetentionProperties") + table = self.timestreamwrite_backend.create_table( + database_name, table_name, retention_properties + ) + return json.dumps(dict(Table=table.description())) + + def delete_table(self): + database_name = self._get_param("DatabaseName") + table_name = self._get_param("TableName") + self.timestreamwrite_backend.delete_table(database_name, table_name) + return "{}" + + def describe_table(self): + database_name = self._get_param("DatabaseName") + table_name = self._get_param("TableName") + table = self.timestreamwrite_backend.describe_table(database_name, table_name) + return json.dumps(dict(Table=table.description())) + + def list_tables(self): + database_name = self._get_param("DatabaseName") + tables = self.timestreamwrite_backend.list_tables(database_name) + return json.dumps(dict(Tables=[t.description() for t in tables])) + + def update_table(self): + database_name = self._get_param("DatabaseName") + table_name = self._get_param("TableName") + retention_properties = self._get_param("RetentionProperties") + table = self.timestreamwrite_backend.update_table( + database_name, table_name, retention_properties + ) + return json.dumps(dict(Table=table.description())) + + def write_records(self): + database_name = self._get_param("DatabaseName") + table_name = self._get_param("TableName") + records = self._get_param("Records") + self.timestreamwrite_backend.write_records(database_name, table_name, records) + return "{}" + + def describe_endpoints(self): + resp = self.timestreamwrite_backend.describe_endpoints() + return json.dumps(resp) diff --git a/moto/timestreamwrite/urls.py b/moto/timestreamwrite/urls.py new file mode 100644 index 000000000..b965cd5b8 --- /dev/null +++ b/moto/timestreamwrite/urls.py @@ -0,0 +1,12 @@ +from .responses import TimestreamWriteResponse + +url_bases = [ + r"https?://ingest\.timestream\.(.+)\.amazonaws\.com", + r"https?://ingest\.timestream\.(.+)\.amazonaws\.com/", +] + +response = TimestreamWriteResponse() + +# Boto3 sends a request to 'https://ingest.timestream.amazonaws.com' +# Which means we need two url_paths - one without slash (for boto3), and one with (for other SDK's/API's) +url_paths = {"{0}$": response.dispatch, "{0}/$": response.dispatch} diff --git a/scripts/scaffold.py b/scripts/scaffold.py index 4816075ef..81f0f0965 100755 --- a/scripts/scaffold.py +++ b/scripts/scaffold.py @@ -34,7 +34,7 @@ import boto3 from moto.core.responses import BaseResponse from moto.core import BaseBackend from inflection import singularize -from .implementation_coverage import get_moto_implementation +from implementation_coverage import get_moto_implementation TEMPLATE_DIR = os.path.join(os.path.dirname(__file__), "./template") diff --git a/scripts/template/lib/models.py.j2 b/scripts/template/lib/models.py.j2 index 0cdb84eb3..6543a1f8b 100644 --- a/scripts/template/lib/models.py.j2 +++ b/scripts/template/lib/models.py.j2 @@ -21,12 +21,12 @@ class {{ service_class }}Backend(BaseBackend): {{ escaped_service }}_backends = {} for available_region in Session().get_available_regions("{{ service }}"): - {{ escaped_service }}_backends[available_region] = {{ service_class }}Backend() + {{ escaped_service }}_backends[available_region] = {{ service_class }}Backend(available_region) for available_region in Session().get_available_regions( "{{ service }}", partition_name="aws-us-gov" ): - {{ escaped_service }}_backends[available_region] = {{ service_class }}Backend() + {{ escaped_service }}_backends[available_region] = {{ service_class }}Backend(available_region) for available_region in Session().get_available_regions( "{{ service }}", partition_name="aws-cn" ): - {{ escaped_service }}_backends[available_region] = {{ service_class }}Backend() + {{ escaped_service }}_backends[available_region] = {{ service_class }}Backend(available_region) diff --git a/scripts/template/test/test_server.py.j2 b/scripts/template/test/test_server.py.j2 index dada73915..805f95aeb 100644 --- a/scripts/template/test/test_server.py.j2 +++ b/scripts/template/test/test_server.py.j2 @@ -1,5 +1,5 @@ """Test different server responses.""" -import sure # noqa +import sure # noqa # pylint: disable=unused-import import moto.server as server from moto import mock_{{ escaped_service }} diff --git a/scripts/template/test/test_service.py.j2 b/scripts/template/test/test_service.py.j2 index 7c1b3b7ed..6308e679d 100644 --- a/scripts/template/test/test_service.py.j2 +++ b/scripts/template/test/test_service.py.j2 @@ -1,7 +1,7 @@ """Unit tests for {{ escaped_service }}-supported APIs.""" import boto3 -import sure # noqa +import sure # noqa # pylint: disable=unused-import from moto import mock_{{ escaped_service }} diff --git a/tests/test_timestreamwrite/__init__.py b/tests/test_timestreamwrite/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/test_timestreamwrite/test_server.py b/tests/test_timestreamwrite/test_server.py new file mode 100644 index 000000000..c1c907268 --- /dev/null +++ b/tests/test_timestreamwrite/test_server.py @@ -0,0 +1,16 @@ +import json +import sure # noqa # pylint: disable=unused-import + +import moto.server as server +from moto import mock_timestreamwrite + + +@mock_timestreamwrite +def test_timestreamwrite_list(): + backend = server.create_backend_app("timestream-write") + test_client = backend.test_client() + + headers = {"X-Amz-Target": "Timestream_20181101.ListDatabases"} + resp = test_client.post("/", headers=headers, json={}) + resp.status_code.should.equal(200) + json.loads(resp.data).should.equal({"Databases": []}) diff --git a/tests/test_timestreamwrite/test_timestreamwrite_database.py b/tests/test_timestreamwrite/test_timestreamwrite_database.py new file mode 100644 index 000000000..0b07565ce --- /dev/null +++ b/tests/test_timestreamwrite/test_timestreamwrite_database.py @@ -0,0 +1,107 @@ +import boto3 +import sure # noqa # pylint: disable=unused-import +from moto import mock_timestreamwrite +from moto.core import ACCOUNT_ID + + +@mock_timestreamwrite +def test_create_database_simple(): + ts = boto3.client("timestream-write", region_name="us-east-1") + resp = ts.create_database(DatabaseName="mydatabase") + database = resp["Database"] + + database.should.have.key("Arn").equals( + f"arn:aws:timestream:us-east-1:{ACCOUNT_ID}:database/mydatabase" + ) + database.should.have.key("DatabaseName").equals("mydatabase") + database.should.have.key("TableCount").equals(0) + database.shouldnt.have.key("KmsKeyId") + + +@mock_timestreamwrite +def test_create_database_advanced(): + ts = boto3.client("timestream-write", region_name="us-east-1") + resp = ts.create_database( + DatabaseName="mydatabase", + KmsKeyId="mykey", + Tags=[{"Key": "k1", "Value": "v1"}, {"Key": "k2", "Value": "v2"}], + ) + database = resp["Database"] + + database.should.have.key("Arn").equals( + f"arn:aws:timestream:us-east-1:{ACCOUNT_ID}:database/mydatabase" + ) + database.should.have.key("DatabaseName").equals("mydatabase") + database.should.have.key("TableCount").equals(0) + database.should.have.key("KmsKeyId").equal("mykey") + + +@mock_timestreamwrite +def test_describe_database(): + ts = boto3.client("timestream-write", region_name="us-east-1") + ts.create_database(DatabaseName="mydatabase", KmsKeyId="mykey") + + database = ts.describe_database(DatabaseName="mydatabase")["Database"] + + database.should.have.key("Arn").equals( + f"arn:aws:timestream:us-east-1:{ACCOUNT_ID}:database/mydatabase" + ) + database.should.have.key("DatabaseName").equals("mydatabase") + database.should.have.key("TableCount").equals(0) + database.should.have.key("KmsKeyId").equal("mykey") + + +@mock_timestreamwrite +def test_list_databases(): + ts = boto3.client("timestream-write", region_name="us-east-1") + ts.create_database(DatabaseName="db_with", KmsKeyId="mykey") + ts.create_database(DatabaseName="db_without") + + resp = ts.list_databases() + databases = resp["Databases"] + databases.should.have.length_of(2) + databases.should.contain( + { + "Arn": f"arn:aws:timestream:us-east-1:{ACCOUNT_ID}:database/db_with", + "DatabaseName": "db_with", + "TableCount": 0, + "KmsKeyId": "mykey", + } + ) + databases.should.contain( + { + "Arn": f"arn:aws:timestream:us-east-1:{ACCOUNT_ID}:database/db_without", + "DatabaseName": "db_without", + "TableCount": 0, + } + ) + + +@mock_timestreamwrite +def test_delete_database(): + ts = boto3.client("timestream-write", region_name="us-east-1") + ts.create_database(DatabaseName="db_1", KmsKeyId="mykey") + ts.create_database(DatabaseName="db_2") + ts.create_database(DatabaseName="db_3", KmsKeyId="mysecondkey") + + ts.list_databases()["Databases"].should.have.length_of(3) + + ts.delete_database(DatabaseName="db_2") + + databases = ts.list_databases()["Databases"] + databases.should.have.length_of(2) + [db["DatabaseName"] for db in databases].should.equal(["db_1", "db_3"]) + + +@mock_timestreamwrite +def test_update_database(): + ts = boto3.client("timestream-write", region_name="us-east-1") + ts.create_database(DatabaseName="mydatabase", KmsKeyId="mykey") + resp = ts.update_database(DatabaseName="mydatabase", KmsKeyId="updatedkey") + resp.should.have.key("Database") + database = resp["Database"] + database.should.have.key("Arn") + database.should.have.key("KmsKeyId").equal("updatedkey") + + database = ts.describe_database(DatabaseName="mydatabase")["Database"] + database.should.have.key("KmsKeyId").equal("updatedkey") diff --git a/tests/test_timestreamwrite/test_timestreamwrite_table.py b/tests/test_timestreamwrite/test_timestreamwrite_table.py new file mode 100644 index 000000000..1b6875b88 --- /dev/null +++ b/tests/test_timestreamwrite/test_timestreamwrite_table.py @@ -0,0 +1,183 @@ +import boto3 +import sure # noqa # pylint: disable=unused-import +from moto import mock_timestreamwrite +from moto.core import ACCOUNT_ID + + +@mock_timestreamwrite +def test_create_table(): + ts = boto3.client("timestream-write", region_name="us-east-1") + ts.create_database(DatabaseName="mydatabase") + + resp = ts.create_table( + DatabaseName="mydatabase", + TableName="mytable", + RetentionProperties={ + "MemoryStoreRetentionPeriodInHours": 7, + "MagneticStoreRetentionPeriodInDays": 42, + }, + ) + table = resp["Table"] + table.should.have.key("Arn").equal( + f"arn:aws:timestream:us-east-1:{ACCOUNT_ID}:database/mydatabase/table/mytable" + ) + table.should.have.key("TableName").equal("mytable") + table.should.have.key("DatabaseName").equal("mydatabase") + table.should.have.key("TableStatus").equal("ACTIVE") + table.should.have.key("RetentionProperties").should.equal( + { + "MemoryStoreRetentionPeriodInHours": 7, + "MagneticStoreRetentionPeriodInDays": 42, + } + ) + + +@mock_timestreamwrite +def test_create_table_without_retention_properties(): + ts = boto3.client("timestream-write", region_name="us-east-1") + ts.create_database(DatabaseName="mydatabase") + + resp = ts.create_table(DatabaseName="mydatabase", TableName="mytable") + table = resp["Table"] + table.should.have.key("Arn").equal( + f"arn:aws:timestream:us-east-1:{ACCOUNT_ID}:database/mydatabase/table/mytable" + ) + table.should.have.key("TableName").equal("mytable") + table.should.have.key("DatabaseName").equal("mydatabase") + table.should.have.key("TableStatus").equal("ACTIVE") + table.shouldnt.have.key("RetentionProperties") + + +@mock_timestreamwrite +def test_describe_table(): + ts = boto3.client("timestream-write", region_name="us-east-1") + ts.create_database(DatabaseName="mydatabase") + + ts.create_table( + DatabaseName="mydatabase", + TableName="mytable", + RetentionProperties={ + "MemoryStoreRetentionPeriodInHours": 10, + "MagneticStoreRetentionPeriodInDays": 12, + }, + ) + + table = ts.describe_table(DatabaseName="mydatabase", TableName="mytable")["Table"] + table.should.have.key("Arn").equal( + f"arn:aws:timestream:us-east-1:{ACCOUNT_ID}:database/mydatabase/table/mytable" + ) + table.should.have.key("TableName").equal("mytable") + table.should.have.key("DatabaseName").equal("mydatabase") + table.should.have.key("TableStatus").equal("ACTIVE") + table.should.have.key("RetentionProperties").should.equal( + { + "MemoryStoreRetentionPeriodInHours": 10, + "MagneticStoreRetentionPeriodInDays": 12, + } + ) + + +@mock_timestreamwrite +def test_create_multiple_tables(): + ts = boto3.client("timestream-write", region_name="us-east-1") + ts.create_database(DatabaseName="mydatabase") + + for idx in range(0, 5): + ts.create_table( + DatabaseName="mydatabase", + TableName=f"mytable_{idx}", + RetentionProperties={ + "MemoryStoreRetentionPeriodInHours": 7, + "MagneticStoreRetentionPeriodInDays": 42, + }, + ) + + database = ts.describe_database(DatabaseName="mydatabase")["Database"] + + database.should.have.key("TableCount").equals(5) + + tables = ts.list_tables(DatabaseName="mydatabase")["Tables"] + tables.should.have.length_of(5) + set([t["DatabaseName"] for t in tables]).should.equal({"mydatabase"}) + set([t["TableName"] for t in tables]).should.equal( + {"mytable_0", "mytable_1", "mytable_2", "mytable_3", "mytable_4"} + ) + set([t["TableStatus"] for t in tables]).should.equal({"ACTIVE"}) + + +@mock_timestreamwrite +def test_delete_table(): + ts = boto3.client("timestream-write", region_name="us-east-1") + ts.create_database(DatabaseName="mydatabase") + + for idx in range(0, 3): + ts.create_table( + DatabaseName="mydatabase", + TableName=f"mytable_{idx}", + RetentionProperties={ + "MemoryStoreRetentionPeriodInHours": 7, + "MagneticStoreRetentionPeriodInDays": 42, + }, + ) + + tables = ts.list_tables(DatabaseName="mydatabase")["Tables"] + tables.should.have.length_of(3) + + ts.delete_table(DatabaseName="mydatabase", TableName="mytable_1") + + tables = ts.list_tables(DatabaseName="mydatabase")["Tables"] + tables.should.have.length_of(2) + set([t["TableName"] for t in tables]).should.equal({"mytable_0", "mytable_2"}) + + +@mock_timestreamwrite +def test_update_table(): + ts = boto3.client("timestream-write", region_name="us-east-1") + ts.create_database(DatabaseName="mydatabase") + ts.create_table(DatabaseName="mydatabase", TableName="mytable") + + resp = ts.update_table( + DatabaseName="mydatabase", + TableName="mytable", + RetentionProperties={ + "MemoryStoreRetentionPeriodInHours": 1, + "MagneticStoreRetentionPeriodInDays": 2, + }, + ) + table = resp["Table"] + table.should.have.key("RetentionProperties").equals( + { + "MagneticStoreRetentionPeriodInDays": 2, + "MemoryStoreRetentionPeriodInHours": 1, + } + ) + + table = ts.describe_table(DatabaseName="mydatabase", TableName="mytable")["Table"] + table.should.have.key("Arn").equal( + f"arn:aws:timestream:us-east-1:{ACCOUNT_ID}:database/mydatabase/table/mytable" + ) + table.should.have.key("TableName").equal("mytable") + table.should.have.key("DatabaseName").equal("mydatabase") + table.should.have.key("TableStatus").equal("ACTIVE") + table.should.have.key("RetentionProperties").equals( + { + "MagneticStoreRetentionPeriodInDays": 2, + "MemoryStoreRetentionPeriodInHours": 1, + } + ) + + +@mock_timestreamwrite +def test_write_records(): + # The query-feature is not available at the moment, + # so there's no way for us to verify writing records is successful + # For now, we'll just send them off into the ether and pray + ts = boto3.client("timestream-write", region_name="us-east-1") + ts.create_database(DatabaseName="mydatabase") + ts.create_table(DatabaseName="mydatabase", TableName="mytable") + + ts.write_records( + DatabaseName="mydatabase", + TableName="mytable", + Records=[{"Dimensions": [], "MeasureName": "mn", "MeasureValue": "mv"}], + )