From a9e1894e6e9bd43b11eb2eb20166833a14718db0 Mon Sep 17 00:00:00 2001
From: Alif Be <alifb@equinor.com>
Date: Thu, 12 Oct 2023 11:49:11 +0200
Subject: [PATCH] Merge rft also on report_step

---
 .../merge_rft_ertobs/merge_rft_ertobs.py      |  4 +-
 tests/test_merge_rft_ertobs.py                | 69 +++++++++++++++++++
 2 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/src/subscript/merge_rft_ertobs/merge_rft_ertobs.py b/src/subscript/merge_rft_ertobs/merge_rft_ertobs.py
index e0a473fe8..078d07aa2 100644
--- a/src/subscript/merge_rft_ertobs/merge_rft_ertobs.py
+++ b/src/subscript/merge_rft_ertobs/merge_rft_ertobs.py
@@ -93,7 +93,7 @@ def split_wellname_reportstep(wellname_reportstep: str) -> Tuple[str, int]:
         A-5_99 gives (A-5_99, 1)  # report steps more than 10 not supported.
         "R_A4_1" gives (R_A4, 1)
         "R_A4" gives (R_A4, 1)
-        "R_A_4" gives (R_A, 4)  # Warning, this is probaly unintended!
+        "R_A_4" gives (R_A, 4)  # Warning, this is probably unintended!
 
     Args:
         wellname_reportstep
@@ -224,7 +224,7 @@ def merge_rft_ertobs(gendatacsv: str, obsdir: str) -> pd.DataFrame:
     # For each simulated well, look up
     logger.info("Parsed %s observations from files in %s", str(len(obs_df)), obsdir)
 
-    return pd.merge(sim_df, obs_df, how="left", on=["well", "order"])
+    return pd.merge(sim_df, obs_df, how="left", on=["well", "order", "report_step"])
 
 
 def main() -> None:
diff --git a/tests/test_merge_rft_ertobs.py b/tests/test_merge_rft_ertobs.py
index 27125d7af..c78cef95e 100644
--- a/tests/test_merge_rft_ertobs.py
+++ b/tests/test_merge_rft_ertobs.py
@@ -138,6 +138,75 @@ def test_merge_drogon_inactive(drogondata):
     )
 
 
+def test_merge_multiple_timesteps(tmp_path):
+    """Check that multiple timesteps is handled properly"""
+    os.chdir(tmp_path)
+    Path("R_A2_1.obs").write_text("299.230  3.000", encoding="utf8")
+    Path("R_A2_2.obs").write_text("289.120  3.000", encoding="utf8")
+    df_gendata = pd.DataFrame(
+        columns=[
+            "order",
+            "utm_x",
+            "utm_y",
+            "measured_depth",
+            "true_vertical_depth",
+            "zone",
+            "pressure",
+            "valid_zone",
+            "is_active",
+            "i",
+            "j",
+            "k",
+            "well",
+            "time",
+            "report_step",
+        ],
+        data=[
+            [
+                0,
+                460994.9,
+                5933813.29,
+                1697.9,
+                1648.9,
+                "Valysar",
+                297.59930419921875,
+                True,
+                True,
+                21,
+                29,
+                2,
+                "R_A2",
+                "2018-03-01",
+                1,
+            ],
+            [
+                0,
+                460994.9,
+                5933813.29,
+                1697.9,
+                1648.9,
+                "Valysar",
+                287.19930419921875,
+                True,
+                True,
+                21,
+                29,
+                2,
+                "R_A2",
+                "2019-03-01",
+                2,
+            ],
+        ],
+    )
+    df_gendata.to_csv("gendata_rft.csv", index=False)
+    dframe = merge_rft_ertobs("gendata_rft.csv", ".")
+    assert not dframe.empty
+    assert {"pressure", "observed", "error", "well", "report_step", "time"}.issubset(
+        dframe.columns
+    )
+    assert len(dframe) == 2
+
+
 def test_extra_obs_file(drogondata):
     """Test that we will not bail on a stray file"""
     Path("rft/FOO.obs").write_text("FOBOBAR", encoding="utf8")