def test_as_of_dates_entity_index(project_storage):
data = {
"entity_id": [1, 2],
"feature_one": [0.5, 0.6],
"feature_two": [0.5, 0.6],
"label": [0, 1],
}
df = pd.DataFrame.from_dict(data)
labels = df.pop("label")
matrix_store = CSVMatrixStore(project_storage, [], "test")
matrix_store.matrix_label_tuple = df, labels
matrix_store.metadata = {"end_time": "2016-01-01", "indices": ["entity_id"], "label_name": "label"}
assert matrix_store.as_of_dates == ["2016-01-01"]
def test_s3_save():
with mock_s3():
client = boto3.client("s3")
client.create_bucket(Bucket="fake-matrix-bucket", ACL="public-read-write")
for example in matrix_stores():
if not isinstance(example, CSVMatrixStore):
continue
project_storage = ProjectStorage("s3://fake-matrix-bucket")
tosave = CSVMatrixStore(project_storage, [], "test")
tosave.metadata = example.metadata
tosave.matrix_label_tuple = example.matrix_label_tuple
tosave.save()
tocheck = CSVMatrixStore(project_storage, [], "test")
assert tocheck.metadata == example.metadata
assert tocheck.design_matrix.to_dict() == example.design_matrix.to_dict()