TimestreamWrite - ensure test parity with AWS (#6971)

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 @@ class TimestreamTable(BaseModel):
         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 @@ class TimestreamDatabase(BaseModel):
         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 @@ class TimestreamDatabase(BaseModel):
             "DatabaseName": self.name,
             "TableCount": len(self.tables.keys()),
             "KmsKeyId": self.kms_key_id,
+            "CreationTime": self.created_on,
+            "LastUpdatedTime": self.updated_on,
         }
 
 

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

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",

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")
+    db_name = "db_" + str(uuid4())[0:6]
+    table_name = "t_" + str(uuid4())[0:6]
 
-    resp = ts.create_table(
-        DatabaseName="mydatabase",
-        TableName="mytable",
-        MagneticStoreWriteProperties={
+    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")
+    db_name = "db_" + str(uuid4())[0:6]
 
-    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."
+    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

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"]
+    db_name = "db_" + str(uuid4())[0:6]
 
-    ts.tag_resource(
-        ResourceARN=db_arn,
-        Tags=[{"Key": "k2", "Value": "v2"}, {"Key": "k3", "Value": "v3"}],
-    )
+    try:
+        db_arn = ts.create_database(
+            DatabaseName=db_name, Tags=[{"Key": "k1", "Value": "v1"}]
+        )["Database"]["Arn"]
 
-    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.tag_resource(
+            ResourceARN=db_arn,
+            Tags=[{"Key": "k2", "Value": "v2"}, {"Key": "k3", "Value": "v3"}],
+        )
 
-    ts.untag_resource(ResourceARN=db_arn, TagKeys=["k2"])
+        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
 
-    resp = ts.list_tags_for_resource(ResourceARN=db_arn)
-    assert resp["Tags"] == [{"Key": "k1", "Value": "v1"}, {"Key": "k3", "Value": "v3"}]
+        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)