From f8febb0daf2e0d7458a10227dd8cb7045ca99f96 Mon Sep 17 00:00:00 2001 From: Bert Blommers Date: Tue, 31 Oct 2023 18:29:04 -0100 Subject: [PATCH] TimestreamWrite - ensure test parity with AWS (#6971) --- moto/timestreamwrite/models.py | 9 +- tests/test_timestreamwrite/__init__.py | 28 +++ .../test_timestreamwrite_database.py | 44 +++-- .../test_timestreamwrite_table.py | 177 +++++++++++------- .../test_timestreamwrite_tagging.py | 53 +++--- 5 files changed, 211 insertions(+), 100 deletions(-) diff --git a/moto/timestreamwrite/models.py b/moto/timestreamwrite/models.py index d3c0044bc025..19669164f92f 100644 --- a/moto/timestreamwrite/models.py +++ b/moto/timestreamwrite/models.py @@ -1,5 +1,6 @@ from typing import Any, Dict, List, Iterable from moto.core import BaseBackend, BackendDict, BaseModel +from moto.core.utils import unix_time from moto.utilities.tagging_service import TaggingService from .exceptions import ResourceNotFound @@ -19,8 +20,8 @@ def __init__( self.name = table_name self.db_name = db_name self.retention_properties = retention_properties or { - "MemoryStoreRetentionPeriodInHours": 123, - "MagneticStoreRetentionPeriodInDays": 123, + "MemoryStoreRetentionPeriodInHours": 6, + "MagneticStoreRetentionPeriodInDays": 73000, } self.magnetic_store_write_properties = magnetic_store_write_properties or {} self.schema = schema or { @@ -75,6 +76,8 @@ def __init__( self.arn = ( f"arn:aws:timestream:{self.region_name}:{account_id}:database/{self.name}" ) + self.created_on = unix_time() + self.updated_on = unix_time() self.tables: Dict[str, TimestreamTable] = dict() def update(self, kms_key_id: str) -> None: @@ -131,6 +134,8 @@ def description(self) -> Dict[str, Any]: "DatabaseName": self.name, "TableCount": len(self.tables.keys()), "KmsKeyId": self.kms_key_id, + "CreationTime": self.created_on, + "LastUpdatedTime": self.updated_on, } diff --git a/tests/test_timestreamwrite/__init__.py b/tests/test_timestreamwrite/__init__.py index e69de29bb2d1..707d83fb011b 100644 --- a/tests/test_timestreamwrite/__init__.py +++ b/tests/test_timestreamwrite/__init__.py @@ -0,0 +1,28 @@ +import os +from functools import wraps +from moto import mock_timestreamwrite, mock_s3, mock_sts + + +def timestreamwrite_aws_verified(func): + """ + Function that is verified to work against AWS. + Can be run against AWS at any time by setting: + MOTO_TEST_ALLOW_AWS_REQUEST=true + + If this environment variable is not set, the function runs in a `mock_timestreamwrite`, `mock_s3` and `mock_sts` context. + + """ + + @wraps(func) + def pagination_wrapper(): + allow_aws_request = ( + os.environ.get("MOTO_TEST_ALLOW_AWS_REQUEST", "false").lower() == "true" + ) + + if allow_aws_request: + return func() + else: + with mock_timestreamwrite(), mock_s3(), mock_sts(): + return func() + + return pagination_wrapper diff --git a/tests/test_timestreamwrite/test_timestreamwrite_database.py b/tests/test_timestreamwrite/test_timestreamwrite_database.py index e49453e9dea6..3427fe1d91e4 100644 --- a/tests/test_timestreamwrite/test_timestreamwrite_database.py +++ b/tests/test_timestreamwrite/test_timestreamwrite_database.py @@ -1,23 +1,39 @@ import boto3 -from botocore.exceptions import ClientError import pytest +from botocore.exceptions import ClientError +from uuid import uuid4 from moto import mock_timestreamwrite from moto.core import DEFAULT_ACCOUNT_ID as ACCOUNT_ID +from . import timestreamwrite_aws_verified -@mock_timestreamwrite + +@pytest.mark.aws_verified +@timestreamwrite_aws_verified def test_create_database_simple(): ts = boto3.client("timestream-write", region_name="us-east-1") - resp = ts.create_database(DatabaseName="mydatabase") - database = resp["Database"] + db_name = "db_" + str(uuid4())[0:6] - assert database["Arn"] == ( - f"arn:aws:timestream:us-east-1:{ACCOUNT_ID}:database/mydatabase" - ) - assert database["DatabaseName"] == "mydatabase" - assert database["TableCount"] == 0 - assert database["KmsKeyId"] == f"arn:aws:kms:us-east-1:{ACCOUNT_ID}:key/default_key" + identity = boto3.client("sts", region_name="us-east-1").get_caller_identity() + account_id = identity["Account"] + + try: + database = ts.create_database(DatabaseName=db_name)["Database"] + + assert ( + database["Arn"] + == f"arn:aws:timestream:us-east-1:{account_id}:database/{db_name}" + ) + assert db_name == db_name + assert database["TableCount"] == 0 + assert database["KmsKeyId"].startswith( + f"arn:aws:kms:us-east-1:{account_id}:key/" + ) + assert "CreationTime" in database + assert "LastUpdatedTime" in database + finally: + ts.delete_database(DatabaseName=db_name) @mock_timestreamwrite @@ -53,7 +69,8 @@ def test_describe_database(): assert database["KmsKeyId"] == "mykey" -@mock_timestreamwrite +@pytest.mark.aws_verified +@timestreamwrite_aws_verified def test_describe_unknown_database(): ts = boto3.client("timestream-write", region_name="us-east-1") with pytest.raises(ClientError) as exc: @@ -72,6 +89,11 @@ def test_list_databases(): resp = ts.list_databases() databases = resp["Databases"] assert len(databases) == 2 + + for db in databases: + db.pop("CreationTime") + db.pop("LastUpdatedTime") + assert { "Arn": f"arn:aws:timestream:us-east-1:{ACCOUNT_ID}:database/db_with", "DatabaseName": "db_with", diff --git a/tests/test_timestreamwrite/test_timestreamwrite_table.py b/tests/test_timestreamwrite/test_timestreamwrite_table.py index 33fa936e59d4..3c2727a86db6 100644 --- a/tests/test_timestreamwrite/test_timestreamwrite_table.py +++ b/tests/test_timestreamwrite/test_timestreamwrite_table.py @@ -1,86 +1,125 @@ +import boto3 +import pytest import time -import boto3 from botocore.exceptions import ClientError -import pytest +from uuid import uuid4 from moto import mock_timestreamwrite, settings from moto.core import DEFAULT_ACCOUNT_ID as ACCOUNT_ID +from . import timestreamwrite_aws_verified -@mock_timestreamwrite + +@pytest.mark.aws_verified +@timestreamwrite_aws_verified def test_create_table(): ts = boto3.client("timestream-write", region_name="us-east-1") - ts.create_database(DatabaseName="mydatabase") + db_name = "db_" + str(uuid4())[0:6] + table_name = "t_" + str(uuid4())[0:6] - resp = ts.create_table( - DatabaseName="mydatabase", - TableName="mytable", - RetentionProperties={ + identity = boto3.client("sts", region_name="us-east-1").get_caller_identity() + account_id = identity["Account"] + + try: + ts.create_database(DatabaseName=db_name) + + resp = ts.create_table( + DatabaseName=db_name, + TableName=table_name, + RetentionProperties={ + "MemoryStoreRetentionPeriodInHours": 7, + "MagneticStoreRetentionPeriodInDays": 42, + }, + ) + table = resp["Table"] + assert table["Arn"] == ( + f"arn:aws:timestream:us-east-1:{account_id}:database/{db_name}/table/{table_name}" + ) + assert table["TableName"] == table_name + assert table["DatabaseName"] == db_name + assert table["TableStatus"] == "ACTIVE" + assert table["RetentionProperties"] == { "MemoryStoreRetentionPeriodInHours": 7, "MagneticStoreRetentionPeriodInDays": 42, - }, - ) - table = resp["Table"] - assert table["Arn"] == ( - f"arn:aws:timestream:us-east-1:{ACCOUNT_ID}:database/mydatabase/table/mytable" - ) - assert table["TableName"] == "mytable" - assert table["DatabaseName"] == "mydatabase" - assert table["TableStatus"] == "ACTIVE" - assert table["RetentionProperties"] == { - "MemoryStoreRetentionPeriodInHours": 7, - "MagneticStoreRetentionPeriodInDays": 42, - } + } + finally: + ts.delete_table(DatabaseName=db_name, TableName=table_name) + ts.delete_database(DatabaseName=db_name) -@mock_timestreamwrite +@pytest.mark.aws_verified +@timestreamwrite_aws_verified def test_create_table__with_magnetic_store_write_properties(): ts = boto3.client("timestream-write", region_name="us-east-1") - ts.create_database(DatabaseName="mydatabase") - - resp = ts.create_table( - DatabaseName="mydatabase", - TableName="mytable", - MagneticStoreWriteProperties={ + db_name = "db_" + str(uuid4())[0:6] + table_name = "t_" + str(uuid4())[0:6] + + identity = boto3.client("sts", region_name="us-east-1").get_caller_identity() + account_id = identity["Account"] + + bucket_name = f"b-{str(uuid4())[0:6]}" + s3 = boto3.client("s3", region_name="us-east-1") + s3.create_bucket(Bucket=bucket_name) + + try: + ts.create_database(DatabaseName=db_name) + + table = ts.create_table( + DatabaseName=db_name, + TableName=table_name, + MagneticStoreWriteProperties={ + "EnableMagneticStoreWrites": True, + "MagneticStoreRejectedDataLocation": { + "S3Configuration": {"BucketName": bucket_name} + }, + }, + )["Table"] + assert table["Arn"] == ( + f"arn:aws:timestream:us-east-1:{account_id}:database/{db_name}/table/{table_name}" + ) + assert table["TableName"] == table_name + assert table["DatabaseName"] == db_name + assert table["TableStatus"] == "ACTIVE" + assert table["MagneticStoreWriteProperties"] == { "EnableMagneticStoreWrites": True, "MagneticStoreRejectedDataLocation": { - "S3Configuration": {"BucketName": "hithere"} + "S3Configuration": {"BucketName": bucket_name} }, - }, - ) - table = resp["Table"] - assert table["Arn"] == ( - f"arn:aws:timestream:us-east-1:{ACCOUNT_ID}:database/mydatabase/table/mytable" - ) - assert table["TableName"] == "mytable" - assert table["DatabaseName"] == "mydatabase" - assert table["TableStatus"] == "ACTIVE" - assert table["MagneticStoreWriteProperties"] == { - "EnableMagneticStoreWrites": True, - "MagneticStoreRejectedDataLocation": { - "S3Configuration": {"BucketName": "hithere"} - }, - } + } + finally: + ts.delete_table(DatabaseName=db_name, TableName=table_name) + ts.delete_database(DatabaseName=db_name) + s3.delete_bucket(Bucket=bucket_name) -@mock_timestreamwrite +@pytest.mark.aws_verified +@timestreamwrite_aws_verified def test_create_table_without_retention_properties(): ts = boto3.client("timestream-write", region_name="us-east-1") - ts.create_database(DatabaseName="mydatabase") + db_name = "db_" + str(uuid4())[0:6] + table_name = "t_" + str(uuid4())[0:6] - resp = ts.create_table(DatabaseName="mydatabase", TableName="mytable") - table = resp["Table"] - assert table["Arn"] == ( - f"arn:aws:timestream:us-east-1:{ACCOUNT_ID}:database/mydatabase/table/mytable" - ) - assert table["TableName"] == "mytable" - assert table["DatabaseName"] == "mydatabase" - assert table["TableStatus"] == "ACTIVE" - assert table["RetentionProperties"] == { - "MemoryStoreRetentionPeriodInHours": 123, - "MagneticStoreRetentionPeriodInDays": 123, - } + identity = boto3.client("sts", region_name="us-east-1").get_caller_identity() + account_id = identity["Account"] + + try: + ts.create_database(DatabaseName=db_name) + + table = ts.create_table(DatabaseName=db_name, TableName=table_name)["Table"] + assert table["Arn"] == ( + f"arn:aws:timestream:us-east-1:{account_id}:database/{db_name}/table/{table_name}" + ) + assert table["TableName"] == table_name + assert table["DatabaseName"] == db_name + assert table["TableStatus"] == "ACTIVE" + assert table["RetentionProperties"] == { + "MemoryStoreRetentionPeriodInHours": 6, + "MagneticStoreRetentionPeriodInDays": 73000, + } + finally: + ts.delete_table(DatabaseName=db_name, TableName=table_name) + ts.delete_database(DatabaseName=db_name) @mock_timestreamwrite @@ -110,16 +149,22 @@ def test_describe_table(): } -@mock_timestreamwrite +@pytest.mark.aws_verified +@timestreamwrite_aws_verified def test_describe_unknown_database(): ts = boto3.client("timestream-write", region_name="us-east-1") - ts.create_database(DatabaseName="mydatabase") - - with pytest.raises(ClientError) as exc: - ts.describe_table(DatabaseName="mydatabase", TableName="unknown") - err = exc.value.response["Error"] - assert err["Code"] == "ResourceNotFoundException" - assert err["Message"] == "The table unknown does not exist." + db_name = "db_" + str(uuid4())[0:6] + + try: + ts.create_database(DatabaseName=db_name) + + with pytest.raises(ClientError) as exc: + ts.describe_table(DatabaseName=db_name, TableName="unknown") + err = exc.value.response["Error"] + assert err["Code"] == "ResourceNotFoundException" + assert err["Message"] == "The table unknown does not exist." + finally: + ts.delete_database(DatabaseName=db_name) @mock_timestreamwrite diff --git a/tests/test_timestreamwrite/test_timestreamwrite_tagging.py b/tests/test_timestreamwrite/test_timestreamwrite_tagging.py index 584c360af4e4..0ad0d178a2df 100644 --- a/tests/test_timestreamwrite/test_timestreamwrite_tagging.py +++ b/tests/test_timestreamwrite/test_timestreamwrite_tagging.py @@ -1,7 +1,11 @@ import boto3 +import pytest +from uuid import uuid4 from moto import mock_timestreamwrite +from . import timestreamwrite_aws_verified + @mock_timestreamwrite def test_list_tagging_for_table_without_tags(): @@ -62,26 +66,33 @@ def test_list_tagging_for_database_with_tags(): assert resp["Tags"] == [{"Key": "k1", "Value": "v1"}] -@mock_timestreamwrite +@pytest.mark.aws_verified +@timestreamwrite_aws_verified def test_tag_and_untag_database(): ts = boto3.client("timestream-write", region_name="us-east-1") - db_arn = ts.create_database( - DatabaseName="mydatabase", Tags=[{"Key": "k1", "Value": "v1"}] - )["Database"]["Arn"] - - ts.tag_resource( - ResourceARN=db_arn, - Tags=[{"Key": "k2", "Value": "v2"}, {"Key": "k3", "Value": "v3"}], - ) - - resp = ts.list_tags_for_resource(ResourceARN=db_arn) - assert resp["Tags"] == [ - {"Key": "k1", "Value": "v1"}, - {"Key": "k2", "Value": "v2"}, - {"Key": "k3", "Value": "v3"}, - ] - - ts.untag_resource(ResourceARN=db_arn, TagKeys=["k2"]) - - resp = ts.list_tags_for_resource(ResourceARN=db_arn) - assert resp["Tags"] == [{"Key": "k1", "Value": "v1"}, {"Key": "k3", "Value": "v3"}] + db_name = "db_" + str(uuid4())[0:6] + + try: + db_arn = ts.create_database( + DatabaseName=db_name, Tags=[{"Key": "k1", "Value": "v1"}] + )["Database"]["Arn"] + + ts.tag_resource( + ResourceARN=db_arn, + Tags=[{"Key": "k2", "Value": "v2"}, {"Key": "k3", "Value": "v3"}], + ) + + tags = ts.list_tags_for_resource(ResourceARN=db_arn)["Tags"] + assert len(tags) == 3 + assert {"Key": "k1", "Value": "v1"} in tags + assert {"Key": "k2", "Value": "v2"} in tags + assert {"Key": "k3", "Value": "v3"} in tags + + ts.untag_resource(ResourceARN=db_arn, TagKeys=["k2"]) + + tags = ts.list_tags_for_resource(ResourceARN=db_arn)["Tags"] + assert len(tags) == 2 + assert {"Key": "k1", "Value": "v1"} in tags + assert {"Key": "k3", "Value": "v3"} in tags + finally: + ts.delete_database(DatabaseName=db_name)