Update Response for TimestreamWrite:WriteRecords (#4950)

The botocore `WriteRecordsResponse` model was updated to include `RecordsIngested` a few months ago.[1]

* Include `RecordsIngested` in the `moto` response
* Use sample records from AWS documentation in `test_write_records`

[1]:4a2fc7f7c0
This commit is contained in:
Brian Pandola 2022-03-18 20:36:22 -07:00 committed by GitHub
parent 0fcf6529ab
commit 411ce71d3a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 51 additions and 18 deletions

View File

@ -85,7 +85,14 @@ class TimestreamWriteResponse(BaseResponse):
table_name = self._get_param("TableName")
records = self._get_param("Records")
self.timestreamwrite_backend.write_records(database_name, table_name, records)
return "{}"
resp = {
"RecordsIngested": {
"Total": len(records),
"MemoryStore": len(records),
"MagneticStore": 0,
}
}
return json.dumps(resp)
def describe_endpoints(self):
resp = self.timestreamwrite_backend.describe_endpoints()

View File

@ -1,3 +1,4 @@
import time
import boto3
import sure # noqa # pylint: disable=unused-import
from moto import mock_timestreamwrite, settings
@ -176,33 +177,58 @@ def test_write_records():
ts.create_database(DatabaseName="mydatabase")
ts.create_table(DatabaseName="mydatabase", TableName="mytable")
ts.write_records(
# Sample records from https://docs.aws.amazon.com/timestream/latest/developerguide/code-samples.write.html
dimensions = [
{"Name": "region", "Value": "us-east-1"},
{"Name": "az", "Value": "az1"},
{"Name": "hostname", "Value": "host1"},
]
cpu_utilization = {
"Dimensions": dimensions,
"MeasureName": "cpu_utilization",
"MeasureValue": "13.5",
"MeasureValueType": "DOUBLE",
"Time": str(time.time()),
}
memory_utilization = {
"Dimensions": dimensions,
"MeasureName": "memory_utilization",
"MeasureValue": "40",
"MeasureValueType": "DOUBLE",
"Time": str(time.time()),
}
sample_records = [cpu_utilization, memory_utilization]
resp = ts.write_records(
DatabaseName="mydatabase",
TableName="mytable",
Records=[{"Dimensions": [], "MeasureName": "mn1", "MeasureValue": "mv1"}],
)
Records=sample_records,
).get("RecordsIngested", {})
resp["Total"].should.equal(len(sample_records))
(resp["MemoryStore"] + resp["MagneticStore"]).should.equal(resp["Total"])
if not settings.TEST_SERVER_MODE:
from moto.timestreamwrite.models import timestreamwrite_backends
backend = timestreamwrite_backends["us-east-1"]
records = backend.databases["mydatabase"].tables["mytable"].records
records.should.equal(
[{"Dimensions": [], "MeasureName": "mn1", "MeasureValue": "mv1"}]
)
records.should.equal(sample_records)
disk_utilization = {
"Dimensions": dimensions,
"MeasureName": "disk_utilization",
"MeasureValue": "100",
"MeasureValueType": "DOUBLE",
"Time": str(time.time()),
}
sample_records.append(disk_utilization)
ts.write_records(
DatabaseName="mydatabase",
TableName="mytable",
Records=[
{"Dimensions": [], "MeasureName": "mn2", "MeasureValue": "mv2"},
{"Dimensions": [], "MeasureName": "mn3", "MeasureValue": "mv3"},
],
)
records.should.equal(
[
{"Dimensions": [], "MeasureName": "mn1", "MeasureValue": "mv1"},
{"Dimensions": [], "MeasureName": "mn2", "MeasureValue": "mv2"},
{"Dimensions": [], "MeasureName": "mn3", "MeasureValue": "mv3"},
]
Records=[disk_utilization],
)
records.should.equal(sample_records)