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") table_name = self._get_param("TableName")
records = self._get_param("Records") records = self._get_param("Records")
self.timestreamwrite_backend.write_records(database_name, table_name, 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): def describe_endpoints(self):
resp = self.timestreamwrite_backend.describe_endpoints() resp = self.timestreamwrite_backend.describe_endpoints()

View File

@ -1,3 +1,4 @@
import time
import boto3 import boto3
import sure # noqa # pylint: disable=unused-import import sure # noqa # pylint: disable=unused-import
from moto import mock_timestreamwrite, settings from moto import mock_timestreamwrite, settings
@ -176,33 +177,58 @@ def test_write_records():
ts.create_database(DatabaseName="mydatabase") ts.create_database(DatabaseName="mydatabase")
ts.create_table(DatabaseName="mydatabase", TableName="mytable") 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", DatabaseName="mydatabase",
TableName="mytable", 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: if not settings.TEST_SERVER_MODE:
from moto.timestreamwrite.models import timestreamwrite_backends from moto.timestreamwrite.models import timestreamwrite_backends
backend = timestreamwrite_backends["us-east-1"] backend = timestreamwrite_backends["us-east-1"]
records = backend.databases["mydatabase"].tables["mytable"].records records = backend.databases["mydatabase"].tables["mytable"].records
records.should.equal( records.should.equal(sample_records)
[{"Dimensions": [], "MeasureName": "mn1", "MeasureValue": "mv1"}]
) disk_utilization = {
"Dimensions": dimensions,
"MeasureName": "disk_utilization",
"MeasureValue": "100",
"MeasureValueType": "DOUBLE",
"Time": str(time.time()),
}
sample_records.append(disk_utilization)
ts.write_records( ts.write_records(
DatabaseName="mydatabase", DatabaseName="mydatabase",
TableName="mytable", TableName="mytable",
Records=[ Records=[disk_utilization],
{"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.should.equal(sample_records)