From 8f23b8e9621d1cb4c335b8eae99112562cd62416 Mon Sep 17 00:00:00 2001
From: DanSava <dan.sava42@gmail.com>
Date: Tue, 21 Dec 2021 12:45:30 +0200
Subject: [PATCH] Add tests for labeled parameters

---
 tests/data/snake_oil_data.py        | 58 +++++++++++++++++++++++++++++
 tests/models/test_ensemble_model.py | 33 ++++++++++++++++
 2 files changed, 91 insertions(+)

diff --git a/tests/data/snake_oil_data.py b/tests/data/snake_oil_data.py
index 1dab4ca3..8d715113 100644
--- a/tests/data/snake_oil_data.py
+++ b/tests/data/snake_oil_data.py
@@ -267,3 +267,61 @@ def to_parquet_helper(dataframe: pd.DataFrame) -> bytes:
         index=["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"],
     ).transpose()
 )
+
+ensembles_response.update({
+    "http://127.0.0.1:5000/ensembles/42": {
+        "data": {
+            "ensemble": {
+                "children": [],
+                "experiment": {"id": "exp1_id"},
+                "parent": None,
+                "id": 1,
+                "timeCreated": "2020-04-29T09:36:26",
+                "size": 1,
+                "activeRealizations": [0],
+                "userdata": '{"name": "default"}',
+            }
+        }
+    },
+    "http://127.0.0.1:5000/ensembles/42/parameters": [
+        "test_parameter_1",
+        "test_parameter_2",
+    ],
+    "http://127.0.0.1:5000/ensembles/42/responses": {
+        "test_resposne": {
+            "name": "name_test_response",
+            "id": "test_response_id_1",
+        },
+    },
+    "http://127.0.0.1:5000/ensembles/42/records/test_parameter_1/labels": [],
+    "http://127.0.0.1:5000/ensembles/42/records/test_parameter_2/labels": ["a", "b"],
+    "http://127.0.0.1:5000/ensembles/42/records/test_parameter_2/labels": ["a", "b"],
+})
+
+ensembles_response[''
+    "http://127.0.0.1:5000/ensembles/42/records/test_parameter_1?"
+] = to_parquet_helper(
+    pd.DataFrame(
+        [0.1, 1.1, 2.1],
+        columns=["0"],
+        index=["0", "1", "2"],
+    ).transpose()
+)
+ensembles_response[
+    "http://127.0.0.1:5000/ensembles/42/records/test_parameter_2?label=a"
+] = to_parquet_helper(
+    pd.DataFrame(
+        [0.01, 1.01, 2.01],
+        columns=["a"],
+        index=["0", "1", "2"],
+    ).transpose()
+)
+ensembles_response[
+    "http://127.0.0.1:5000/ensembles/42/records/test_parameter_2?label=b"
+] = to_parquet_helper(
+    pd.DataFrame(
+        [0.02, 1.02, 2.02],
+        columns=["b"],
+        index=["0", "1", "2"],
+    ).transpose()
+)
diff --git a/tests/models/test_ensemble_model.py b/tests/models/test_ensemble_model.py
index 803a3d4e..caae1327 100644
--- a/tests/models/test_ensemble_model.py
+++ b/tests/models/test_ensemble_model.py
@@ -7,3 +7,36 @@ def test_ensemble_model(mock_data):
     assert ens_model.children[0]._name == "default_smoother_update"
     assert ens_model._name == "default"
     assert len(ens_model.responses) == 1
+
+def test_ensemble_model_labled_parameters(mock_data):
+    ens_id = 42
+    ens_model = EnsembleModel(ensemble_id=ens_id, project_id=None)
+    assert ens_model._name == "default"
+    assert len(ens_model.parameters) == 3
+    for param_name, parameter in ens_model.parameters.items():
+        name, label = param_name.split("::", maxsplit=1) if "::" in param_name else [param_name, None]
+        expected_lables = ens_model._data_loader.get_record_labels(ens_id, name)
+        if label is not None:
+            assert label in expected_lables
+
+def test_ensemble_model_parameter_data(mock_data):
+    ens_id = 42
+    ens_model = EnsembleModel(ensemble_id=ens_id, project_id=None)
+    parameters = ens_model.parameters
+    assert len(parameters) == 3
+
+    # Parameter no lables:
+    expected_lables = ens_model._data_loader.get_record_labels(ens_id, "test_parameter_1")
+    assert expected_lables == []
+    data = parameters["test_parameter_1"].data_df().values
+    assert data.flatten().tolist() == [0.1, 1.1, 2.1]
+
+    # Parameter with lables:
+    expected_lables = ens_model._data_loader.get_record_labels(ens_id, "test_parameter_2")
+    assert expected_lables == ["a", "b"]
+    data = parameters["test_parameter_2::a"].data_df()["a"].values.tolist()
+    assert data == [0.01, 1.01, 2.01]
+    data = parameters["test_parameter_2::b"].data_df()["b"].values.tolist()
+    assert data == [0.02, 1.02, 2.02]
+
+