pastastore 1.11.0__py3-none-any.whl → 1.12.1__py3-none-any.whl

tests/conftest.py

@@ -1,5 +1,6 @@
 # ruff: noqa: D100 D103
 import importlib
+import inspect
 from importlib import metadata
 
 import pandas as pd
@@ -8,50 +9,94 @@ import pytest
 
 import pastastore as pst
 
-params = ["dict", "pas", "arcticdb"]
-
-
-def initialize_project(conn):
-    pstore = pst.PastaStore(conn, "test_project")
-
-    # oseries 1
-    o = pd.read_csv("./tests/data/obs.csv", index_col=0, parse_dates=True)
-    pstore.add_oseries(o, "oseries1", metadata={"x": 165000, "y": 424000})
-
-    # oseries 2
-    o = pd.read_csv("./tests/data/head_nb1.csv", index_col=0, parse_dates=True)
-    pstore.add_oseries(o, "oseries2", metadata={"x": 164000, "y": 423000})
+params = [
+    "dict",
+    "pas",
+    "arcticdb",
+]
+
+
+@pytest.fixture(scope="module")
+def data1():
+    d = {
+        "oseries1": pd.read_csv("./tests/data/obs.csv", index_col=0, parse_dates=True),
+        "oseries1_meta": {"x": 165000, "y": 424000},
+        "prec1": pd.read_csv("./tests/data/rain.csv", index_col=0, parse_dates=True),
+        "prec1_meta": {"x": 165050, "y": 424050},
+        "evap1": pd.read_csv("./tests/data/evap.csv", index_col=0, parse_dates=True),
+        "evap1_meta": {"x": 164500, "y": 424000},
+    }
+    return d
+
+
+@pytest.fixture(scope="module")
+def data2():
+    d = {
+        "oseries2": pd.read_csv(
+            "./tests/data/head_nb1.csv", index_col=0, parse_dates=True
+        ),
+        "oseries2_meta": {"x": 164000, "y": 423000},
+        "prec2": pd.read_csv(
+            "./tests/data/rain_nb1.csv", index_col=0, parse_dates=True
+        ),
+        "prec2_meta": {"x": 164010, "y": 423000},
+        "evap2": pd.read_csv(
+            "./tests/data/evap_nb1.csv", index_col=0, parse_dates=True
+        ),
+        "evap2_meta": {"x": 164000, "y": 423030},
+    }
+    return d
+
+
+@pytest.fixture(scope="module")
+def data3():
+    w = pd.read_csv("./tests/data/well_month_end.csv", index_col=0, parse_dates=True)
+    w = ps.ts.timestep_weighted_resample(
+        w,
+        pd.date_range(w.index[0] - pd.offsets.MonthBegin(), w.index[-1], freq="D"),
+    ).bfill()
 
-    # oseries 3
-    o = pd.read_csv("./tests/data/gw_obs.csv", index_col=0, parse_dates=True)
-    pstore.add_oseries(o, "oseries3", metadata={"x": 165554, "y": 422685})
+    d = {
+        "oseries3": pd.read_csv(
+            "./tests/data/gw_obs.csv", index_col=0, parse_dates=True
+        ),
+        "oseries3_meta": {"x": 165554, "y": 422685},
+        "well1": w,
+        "well1_meta": {"x": 164691, "y": 423579},
+        "well2": w + 10,
+        "well2_meta": {"x": 164691 + 2000, "y": 423579 + 2000},  # far away
+    }
+    return d
 
-    # prec 1
-    s = pd.read_csv("./tests/data/rain.csv", index_col=0, parse_dates=True)
-    pstore.add_stress(s, "prec1", kind="prec", metadata={"x": 165050, "y": 424050})
 
-    # prec 2
-    s = pd.read_csv("./tests/data/rain_nb1.csv", index_col=0, parse_dates=True)
-    pstore.add_stress(s, "prec2", kind="prec", metadata={"x": 164010, "y": 423000})
+def initialize_project(conn, data1, data2, data3):
+    pstore = pst.PastaStore(conn, "test_project")
 
-    # evap 1
-    s = pd.read_csv("./tests/data/evap.csv", index_col=0, parse_dates=True)
-    pstore.add_stress(s, "evap1", kind="evap", metadata={"x": 164500, "y": 424000})
+    # dataset 1
+    pstore.add_oseries(data1["oseries1"], "oseries1", metadata=data1["oseries1_meta"])
+    pstore.add_stress(
+        data1["prec1"], "prec1", kind="prec", metadata=data1["prec1_meta"]
+    )
+    pstore.add_stress(
+        data1["evap1"], "evap1", kind="evap", metadata=data1["evap1_meta"]
+    )
 
-    # evap 2
-    s = pd.read_csv("./tests/data/evap_nb1.csv", index_col=0, parse_dates=True)
-    pstore.add_stress(s, "evap2", kind="evap", metadata={"x": 164000, "y": 423030})
+    # dataset 2
+    pstore.add_oseries(data2["oseries2"], "oseries2", metadata=data2["oseries2_meta"])
+    pstore.add_stress(
+        data2["prec2"], "prec2", kind="prec", metadata=data2["prec2_meta"]
+    )
+    pstore.add_stress(
+        data2["evap2"], "evap2", kind="evap", metadata=data2["evap2_meta"]
+    )
 
-    # well 1
-    s = pd.read_csv("./tests/data/well_month_end.csv", index_col=0, parse_dates=True)
-    s = ps.ts.timestep_weighted_resample(
-        s,
-        pd.date_range(s.index[0] - pd.offsets.MonthBegin(), s.index[-1], freq="D"),
-    ).bfill()
-    pstore.add_stress(s, "well1", kind="well", metadata={"x": 164691, "y": 423579})
-    # add second well
+    # dataset 3
+    pstore.add_oseries(data3["oseries3"], "oseries3", metadata=data3["oseries3_meta"])
+    pstore.add_stress(
+        data3["well1"], "well1", kind="well", metadata=data3["well1_meta"]
+    )
     pstore.add_stress(
-        s + 10, "well2", kind="well", metadata={"x": 164691 + 200, "y": 423579_200}
+        data3["well2"], "well2", kind="well", metadata=data3["well2_meta"]
     )
 
     return pstore
@@ -63,7 +108,7 @@ def conn(request):
     name = f"test_{request.param}"
     # connect to dbase
     if request.param == "arcticdb":
-        uri = "lmdb://./arctic_db/"
+        uri = "lmdb://./tests/data/arcticdb/"
         conn = pst.ArcticDBConnector(name, uri)
     elif request.param == "dict":
         conn = pst.DictConnector(name)
@@ -76,32 +121,25 @@ def conn(request):
 
 
 @pytest.fixture(scope="module", params=params)
-def pstore(request):
+def pstore(request, data1, data2, data3):
     if request.param == "arcticdb":
-        name = "test_project"
-        uri = "lmdb://./arctic_db/"
+        name = "testdb"
+        uri = "lmdb://./tests/data/arcticdb/"
         connector = pst.ArcticDBConnector(name, uri)
     elif request.param == "dict":
-        name = "test_project"
+        name = "testdb"
         connector = pst.DictConnector(name)
     elif request.param == "pas":
-        name = "test_project"
+        name = "testdb"
         connector = pst.PasConnector(name, "./tests/data/pas")
     else:
         raise ValueError("Unrecognized parameter!")
-    pstore = initialize_project(connector)
+    pstore = initialize_project(connector, data1, data2, data3)
     pstore.type = request.param  # added here for defining test dependencies
     yield pstore
     pst.util.delete_pastastore(pstore)
 
 
-def delete_arcticdb_test_db():
-    connstr = "lmdb://./arctic_db/"
-    name = "test_project"
-    connector = pst.ArcticDBConnector(name, connstr)
-    pst.util.delete_arcticdb_connector(connector)
-
-
 _has_pkg_cache = {}
 
 
@@ -158,3 +196,67 @@ def requires_pkg(*pkgs):
         reason=f"missing package{'s' if len(missing) != 1 else ''}: "
         + ", ".join(missing),
     )
+
+
+def for_connectors(connectors=None):
+    """Decorate to run tests only for specified connector types.
+
+    Parameters
+    ----------
+    connectors : list of str, optional
+        List of connector types for which the test should run.
+        If None, defaults to all connector types: ["dict", "pas", "arcticdb"].
+
+    Returns
+    -------
+    function
+        Decorated test function that runs only for specified connectors.
+    """
+    if connectors is None:
+        connectors = ["dict", "pas", "arcticdb"]
+
+    def decorator(func):
+        def wrapper(self, *args, **kwargs):
+            # Get the connector type from fixtures passed
+            # Could be in args or kwargs depending on how pytest passes it
+            store = None
+
+            # Check kwargs first (fixture passed by name)
+            for fixture_name in ["pstore_with_models", "pstore", "conn"]:
+                if fixture_name in kwargs:
+                    store = kwargs[fixture_name]
+                    break
+
+            # If not in kwargs, check first positional arg
+            if store is None and len(args) > 0:
+                store = args[0]
+
+            if store is None:
+                # Call without checking (shouldn't happen)
+                return func(self, *args, **kwargs)
+
+            # Check if we should skip based on connector type
+            if hasattr(store, "conn"):
+                conn_type = store.conn.conn_type
+            elif hasattr(store, "conn_type"):
+                conn_type = store.conn_type
+            else:
+                # If we can't determine connector type, just run the test
+                return func(self, *args, **kwargs)
+
+            if conn_type not in connectors:
+                pytest.skip(
+                    f"Test skipped for {conn_type} connector "
+                    f"(only runs for: {connectors})"
+                )
+
+            return func(self, *args, **kwargs)
+
+        # Preserve the original function signature and metadata for pytest
+        wrapper.__signature__ = inspect.signature(func)
+        wrapper.__name__ = func.__name__
+        wrapper.__doc__ = func.__doc__
+
+        return wrapper
+
+    return decorator

tests/test_003_pastastore.py

@@ -39,9 +39,11 @@ def test_get_tmintmax(pstore):
 
 @pytest.mark.dependency
 def test_search(pstore):
-    results = pstore.search("oseries", "OSER", case_sensitive=False)
+    results = pstore.search("OSER", libname="oseries", case_sensitive=False)
     assert len(results) == 3
     assert len(set(results) - {"oseries1", "oseries2", "oseries3"}) == 0
+    results = pstore.search("oser", libname=None, case_sensitive=True)
+    assert len(results["oseries"]) == 3
 
 
 @pytest.mark.dependency

tests/test_009_parallel.py

@@ -0,0 +1,393 @@
+# ruff: noqa: D100 D103
+import pytest
+from conftest import for_connectors
+
+import pastastore as pst
+
+pst.get_color_logger("DEBUG", logger_name="pastastore")
+
+
+@pytest.fixture(scope="module")
+def pstore_with_models(pstore):
+    """
+    Fixture that creates models for oseries1 and oseries2.
+
+    This provides a pastastore with two models ready for testing parallel-safe
+    operations and time series link updates.
+    """
+    # Create and add model for oseries1
+    ml1 = pstore.create_model("oseries1")
+    ml1.solve(report=False)
+    pstore.add_model(ml1)
+
+    # Create and add model for oseries2
+    ml2 = pstore.create_model("oseries2")
+    ml2.solve(report=False)
+    pstore.add_model(ml2)
+    # trigger update to start fresh
+    _ = pstore.conn.oseries_models
+    # ensure clean start
+    pstore.conn._trigger_links_update_if_needed()
+    yield pstore
+
+    # Cleanup
+    if "oseries1" in pstore.model_names:
+        pstore.del_model("oseries1")
+    if "oseries2" in pstore.model_names:
+        pstore.del_model("oseries2")
+
+
+class TestSingleThreadOperations:
+    """Test single-thread add_model operations for all connector types."""
+
+    def test_oseries_models_triggers_update(self, pstore_with_models):
+        """Test that accessing oseries_models triggers update of added models."""
+        # Initially, _added_models should be empty
+        assert len(pstore_with_models.conn._added_models) == 0
+        # add model
+        ml = pstore_with_models.create_model("oseries3")
+        pstore_with_models.add_model(ml)
+        # Access oseries_models to trigger update
+        _ = pstore_with_models.conn.oseries_models
+        # After update, _added_models should be empty
+        assert len(pstore_with_models.conn._added_models) == 0
+        # cleanup
+        pstore_with_models.del_model("oseries3")
+
+    def test_stresses_models_triggers_update(self, pstore_with_models):
+        """Test that accessing stresses_models triggers update of added models."""
+        # initially, _added_models should be empty
+        assert len(pstore_with_models.conn._added_models) == 0
+        # add model
+        ml = pstore_with_models.create_model("oseries3")
+        pstore_with_models.add_model(ml)
+        # access stresses_models to trigger update
+        _ = pstore_with_models.conn.stresses_models
+        # after update, _added_models should be empty
+        assert len(pstore_with_models.conn._added_models) == 0
+        # cleanup
+        pstore_with_models.del_model("oseries3")
+
+    def test_oseries_with_models_triggers_update(self, pstore_with_models):
+        """Test that accessing oseries_with_models triggers update of added models."""
+        # initially, _added_models should be empty
+        assert len(pstore_with_models.conn._added_models) == 0
+        # add model
+        ml = pstore_with_models.create_model("oseries3")
+        pstore_with_models.add_model(ml)
+        # access oseries_models to trigger update
+        _ = pstore_with_models.conn.oseries_with_models
+        # after update, _added_models should be empty
+        assert len(pstore_with_models.conn._added_models) == 0
+        # cleanup
+        pstore_with_models.del_model("oseries3")
+
+    def test_stresses_with_models_triggers_update(self, pstore_with_models):
+        """Test that accessing stresses_with_models triggers update of added models."""
+        # initially, _added_models should be empty
+        assert len(pstore_with_models.conn._added_models) == 0
+        # add model
+        ml = pstore_with_models.create_model("oseries3")
+        pstore_with_models.add_model(ml)
+        # access stresses_with_models to trigger update
+        _ = pstore_with_models.conn.stresses_with_models
+        # after update, _added_models should be empty
+        assert len(pstore_with_models.conn._added_models) == 0
+        # cleanup
+        pstore_with_models.del_model("oseries3")
+
+    def test_add_model_appends_to_added_models_list(self, pstore_with_models):
+        """Test that add_model appends to _added_models list."""
+        # Check internal state, should contain two models
+        assert len(pstore_with_models.conn._added_models) == 0
+        # Create a new model
+        ml = pstore_with_models.create_model("oseries3")
+        # Add model
+        pstore_with_models.add_model(ml)
+        # Verify the model name was added to _added_models
+        assert "oseries3" in pstore_with_models.conn._added_models
+        # check oseries_models (triggers link update)
+        assert "oseries3" in pstore_with_models.oseries_models
+        # check that _added_models is now empty after update
+        assert len(pstore_with_models.conn._added_models) == 0
+        # Cleanup
+        pstore_with_models.del_model("oseries3")
+
+    def test_del_model_deletes_from_added_models_list(self, pstore_with_models):
+        """Test that del_model deletes from non-empty _added_models list."""
+        # Check internal state, should be empty
+        assert len(pstore_with_models.conn._added_models) == 0
+        # Create a new model
+        ml = pstore_with_models.create_model("oseries3")
+        # Add model
+        pstore_with_models.add_model(ml)
+        pstore_with_models.del_model("oseries3")
+        # check that _added_models is now empty after update
+        assert len(pstore_with_models.conn._added_models) == 0
+
+
+class TestParallelButNotReally:
+    """Test internal update flag behavior (for parallel)."""
+
+    @for_connectors(connectors=["pas", "arcticdb"])
+    def test_parallel_add_model_sets_update_flags(self, pstore_with_models):
+        """Test parallel_safe adds set update flags for parallel connectors."""
+
+        def add_model(name):
+            """Add model using global pstore."""
+            ml = pstore_with_models.get_models(name)
+            ml.solve(report=False)
+            pstore_with_models.add_model(ml, overwrite=True)
+
+        pstore_with_models.apply(
+            func=add_model,
+            names=["oseries1", "oseries2"],
+            libname="models",
+            parallel=False,
+        )
+
+        # check update parameters are set to True
+        assert pstore_with_models.conn._oseries_links_need_update.value is True
+        assert pstore_with_models.conn._stresses_links_need_update.value is True
+
+        # trigger update
+        pstore_with_models.conn._trigger_links_update_if_needed()
+
+        # after update, flags should be False
+        assert pstore_with_models.conn._oseries_links_need_update.value is False
+        assert pstore_with_models.conn._stresses_links_need_update.value is False
+
+    @for_connectors(connectors=["pas", "arcticdb"])
+    def test_oseries_models_triggers_update(self, pstore_with_models):
+        """Test that accessing oseries_models triggers update of added models."""
+        # Initially, _added_models should be empty
+        assert len(pstore_with_models.conn._added_models) == 0
+
+        def add_model(name):
+            """Add model using pstore."""
+            ml = pstore_with_models.create_model(name)
+            pstore_with_models.add_model(ml, overwrite=True)
+
+        pstore_with_models.apply(
+            func=add_model,
+            names=["oseries1", "oseries2"],
+            libname="oseries",
+            parallel=False,
+        )
+
+        # check update parameters are set to True
+        assert pstore_with_models.conn._oseries_links_need_update.value is True
+        assert pstore_with_models.conn._stresses_links_need_update.value is True
+
+        # trigger update
+        om = pstore_with_models.oseries_models
+
+        # check if result is correct
+        assert om["oseries1"] == ["oseries1"]
+        assert om["oseries2"] == ["oseries2"]
+
+        # after update, flags should be False
+        assert pstore_with_models.conn._oseries_links_need_update.value is False
+        assert pstore_with_models.conn._stresses_links_need_update.value is False
+
+    @for_connectors(connectors=["pas", "arcticdb"])
+    def test_stresses_models_triggers_update(self, pstore_with_models):
+        """Test that accessing stresses_models triggers update of added models."""
+        # Initially, _added_models should be empty
+        assert len(pstore_with_models.conn._added_models) == 0
+
+        def add_model(name):
+            """Add model using pstore."""
+            ml = pstore_with_models.create_model(name)
+            pstore_with_models.add_model(ml, overwrite=True)
+
+        pstore_with_models.apply(
+            func=add_model,
+            names=["oseries1", "oseries2"],
+            libname="oseries",
+            parallel=False,
+        )
+
+        # check update parameters are set to True
+        assert pstore_with_models.conn._oseries_links_need_update.value is True
+        assert pstore_with_models.conn._stresses_links_need_update.value is True
+
+        # trigger update
+        sm = pstore_with_models.stresses_models
+
+        # check if result is correct
+        assert sm["prec1"] == ["oseries1"]
+        assert sm["prec2"] == ["oseries2"]
+        assert sm["evap1"] == ["oseries1"]
+        assert sm["evap2"] == ["oseries2"]
+
+        # after update, flags should be False
+        assert pstore_with_models.conn._oseries_links_need_update.value is False
+        assert pstore_with_models.conn._stresses_links_need_update.value is False
+
+    @for_connectors(connectors=["pas", "arcticdb"])
+    def test_oseries_with_models_triggers_update(self, pstore_with_models):
+        """Test that accessing oseries_with_models triggers update of added models."""
+        # Initially, _added_models should be empty
+        assert len(pstore_with_models.conn._added_models) == 0
+
+        def add_model(name):
+            """Add model using pstore."""
+            ml = pstore_with_models.create_model(name)
+            pstore_with_models.add_model(ml, overwrite=True)
+
+        pstore_with_models.apply(
+            func=add_model,
+            names=["oseries1", "oseries2"],
+            libname="oseries",
+            parallel=False,
+        )
+
+        # check update parameters are set to True
+        assert pstore_with_models.conn._oseries_links_need_update.value is True
+        assert pstore_with_models.conn._stresses_links_need_update.value is True
+
+        # trigger update
+        owm = pstore_with_models.oseries_with_models
+
+        # check if result is correct
+        assert "oseries1" in owm
+        assert "oseries2" in owm
+
+        # after update, flags should be False
+        assert pstore_with_models.conn._oseries_links_need_update.value is False
+        assert pstore_with_models.conn._stresses_links_need_update.value is False
+
+    @for_connectors(connectors=["pas", "arcticdb"])
+    def test_stresses_with_models_triggers_update(self, pstore_with_models):
+        """Test that accessing stresses_with_models triggers update of added models."""
+        # Initially, _added_models should be empty
+        assert len(pstore_with_models.conn._added_models) == 0
+
+        def add_model(name):
+            """Add model using pstore."""
+            ml = pstore_with_models.create_model(name)
+            pstore_with_models.add_model(ml, overwrite=True)
+
+        pstore_with_models.apply(
+            func=add_model,
+            names=["oseries1", "oseries2"],
+            libname="oseries",
+            parallel=False,
+        )
+
+        # check update parameters are set to True
+        assert pstore_with_models.conn._oseries_links_need_update.value is True
+        assert pstore_with_models.conn._stresses_links_need_update.value is True
+
+        # check if result is correct
+        swm = pstore_with_models.stresses_with_models
+
+        assert "prec1" in swm
+        assert "prec2" in swm
+        assert "evap1" in swm
+        assert "evap2" in swm
+
+        # after update, flags should be False
+        assert pstore_with_models.conn._oseries_links_need_update.value is False
+        assert pstore_with_models.conn._stresses_links_need_update.value is False
+
+
+def add_model_pas(name):
+    """Add model using pstore."""
+    ml = ppstore.create_model(name)
+    ppstore.add_model(ml, overwrite=True)
+
+
+def add_model_arcticdb(name):
+    """Add model using pstore."""
+    ppstore = pst.PastaStore(conn)
+    ml = ppstore.create_model(name)
+    conn.add_model(ml, overwrite=True)
+
+
+def setup_parallel_pstore(conn_type, data1, data2):
+    global conn
+    global ppstore
+
+    if conn_type == "arcticdb":
+        name = "paralleldb"
+        uri = "lmdb://./tests/data/arcticdb/"
+        conn = pst.ArcticDBConnector(name, uri)
+    elif conn_type == "pas":
+        name = "paralleldb"
+        conn = pst.PasConnector(name, "./tests/data/pas")
+    else:
+        raise ValueError("Unrecognized parameter!")
+
+    ppstore = pst.PastaStore(conn)
+    # dataset 1
+    ppstore.add_oseries(data1["oseries1"], "oseries1", metadata=data1["oseries1_meta"])
+    ppstore.add_stress(
+        data1["prec1"], "prec1", kind="prec", metadata=data1["prec1_meta"]
+    )
+    ppstore.add_stress(
+        data1["evap1"], "evap1", kind="evap", metadata=data1["evap1_meta"]
+    )
+
+    # dataset 2
+    ppstore.add_oseries(data2["oseries2"], "oseries2", metadata=data2["oseries2_meta"])
+    ppstore.add_stress(
+        data2["prec2"], "prec2", kind="prec", metadata=data2["prec2_meta"]
+    )
+    ppstore.add_stress(
+        data2["evap2"], "evap2", kind="evap", metadata=data2["evap2_meta"]
+    )
+    return ppstore
+
+
+@pytest.mark.parametrize("conn_type", ["pas", "arcticdb"])
+def test_parallel_add_model(conn_type, data1, data2):
+    """Test that accessing stresses_with_models triggers update of added models."""
+    ppstore = setup_parallel_pstore(conn_type, data1, data2)
+
+    if conn_type == "arcticdb":
+        func = add_model_arcticdb
+    elif conn_type == "pas":
+        func = add_model_pas
+    else:
+        raise ValueError("Unrecognized parameter!")
+    try:
+        # Initially, _added_models should be empty
+        assert len(ppstore.conn._added_models) == 0
+
+        ppstore.apply(
+            func=func,
+            names=["oseries1", "oseries2"],
+            libname="oseries",
+            parallel=True,
+            max_workers=1,
+        )
+
+        # check update parameters are set to False, since after parallel these are
+        # recomputed automatically
+        assert ppstore.conn._oseries_links_need_update.value is False
+        assert ppstore.conn._stresses_links_need_update.value is False
+
+        # check if result is correct
+        om = ppstore.oseries_models
+        owm = ppstore.oseries_with_models
+        sm = ppstore.stresses_models
+        swm = ppstore.stresses_with_models
+        assert "oseries1" in owm
+        assert "oseries2" in owm
+        assert "oseries1" in om
+        assert "oseries2" in om
+        assert "oseries1" in om["oseries1"]
+        assert "oseries2" in om["oseries2"]
+        assert "prec1" in swm
+        assert "prec2" in swm
+        assert "evap1" in swm
+        assert "evap2" in swm
+        assert "oseries1" in sm["prec1"]
+        assert "oseries2" in sm["prec2"]
+        assert "oseries1" in sm["evap1"]
+        assert "oseries2" in sm["evap2"]
+
+    finally:
+        pst.util.delete_pastastore(ppstore)