From 24462713fcb258f6722b954d6b901d6e8ec41355 Mon Sep 17 00:00:00 2001
From: Paul Natsuo Kishimoto <mail@paul.kishimoto.name>
Date: Wed, 13 Sep 2023 14:37:15 +0200
Subject: [PATCH] Add .project.ssp.data with a ExoDataSource for SSP data

---
 message_ix_models/project/ssp/data.py       | 47 +++++++++++++++++++++
 message_ix_models/tests/project/test_ssp.py | 37 ++++++++++++++++
 2 files changed, 84 insertions(+)
 create mode 100644 message_ix_models/project/ssp/data.py

diff --git a/message_ix_models/project/ssp/data.py b/message_ix_models/project/ssp/data.py
new file mode 100644
index 0000000000..8f4a81ae26
--- /dev/null
+++ b/message_ix_models/project/ssp/data.py
@@ -0,0 +1,47 @@
+from message_ix_models.tools.exo_data import (
+    ExoDataSource,
+    iamc_like_data_for_query,
+    register_source,
+)
+from message_ix_models.util import private_data_path
+
+
+@register_source
+class SSPUpdate(ExoDataSource):
+    """Provider of exogenous data from the SSP Update database."""
+
+    id = "SSP update"
+
+    def __init__(self, source, source_kw):
+        s = "ICONICS:SSP(2024)."
+        if not source.startswith(s):
+            raise ValueError(source)
+
+        *parts, self.ssp_number = source.partition(s)
+
+        # Map the `measure` keyword to a string appearing in the data
+        self.measure = {
+            "GDP": "GDP|PPP",
+            "POP": "Population",
+        }[source_kw.pop("measure")]
+
+        # Store the model ID, if any
+        self.model = source_kw.pop("model", None)
+
+        if len(source_kw):
+            raise ValueError(source_kw)
+
+    def __call__(self):
+        # Assemble a query string
+        query = " and ".join(
+            [
+                f"Scenario == 'SSP{self.ssp_number} - Review Phase 1'",
+                f"Variable == '{self.measure}'",
+                f"Model == '{self.model}'" if self.model else "True",
+            ]
+        )
+
+        path = private_data_path("ssp", "SSP-Review-Phase-1.csv.gz")
+        assert path.exists(), "TODO handle the case where message_data is not insalled"
+
+        return iamc_like_data_for_query(path, query)
diff --git a/message_ix_models/tests/project/test_ssp.py b/message_ix_models/tests/project/test_ssp.py
index 75b77c1377..35c1fe473e 100644
--- a/message_ix_models/tests/project/test_ssp.py
+++ b/message_ix_models/tests/project/test_ssp.py
@@ -1,4 +1,5 @@
 import pytest
+from genno import Computer
 
 from message_ix_models.project.ssp import (
     SSP,
@@ -8,6 +9,8 @@
     parse,
     ssp_field,
 )
+from message_ix_models.project.ssp.data import SSPUpdate  # noqa: F401
+from message_ix_models.tools.exo_data import prepare_computer
 
 
 def test_generate(tmp_path, test_context):
@@ -81,3 +84,37 @@ class Foo:
 
 def test_cli(mix_models_cli):
     mix_models_cli.assert_exit_0(["ssp", "gen-structures", "--dry-run"])
+
+
+class TestSSPUpdate:
+    @pytest.mark.parametrize(
+        "source",
+        (
+            "ICONICS:SSP(2024).1",
+            "ICONICS:SSP(2024).2",
+            "ICONICS:SSP(2024).3",
+            "ICONICS:SSP(2024).4",
+            "ICONICS:SSP(2024).5",
+        ),
+    )
+    @pytest.mark.parametrize(
+        "source_kw",
+        (
+            dict(measure="POP"),
+            dict(measure="GDP", model="IIASA GDP 2023"),
+            dict(measure="GDP", model="OECD ENV-Growth 2023"),
+        ),
+    )
+    def test_prepare_computer(self, test_context, source, source_kw):
+        c = Computer()
+        keys = prepare_computer(test_context, c, source, source_kw)
+
+        # Preparation of data runs successfully
+        result = c.get(keys[0])
+
+        # Data has the expected dimensions
+        assert ("n", "y") == result.dims
+
+        # Data is complete
+        assert 14 == len(result.coords["n"])
+        assert 14 == len(result.coords["y"])