Fix #1006

src/ttsim/interface_dag_elements/fail_if.py

@@ -638,7 +638,11 @@ def input_df_mapper_has_incorrect_format(
         raise TypeError(msg)
 
 
-@fail_function()
+@fail_function(
+    include_if_any_element_present=[
+        "raw_results__columns",
+    ]
+)
 def root_nodes_are_missing(
     specialized_environment__tax_transfer_dag: nx.DiGraph,
     specialized_environment__with_partialled_params_and_scalars: SpecEnvWithPartialledParamsAndScalars,

tests/ttsim/interface_dag_elements/test_failures.py

@@ -1365,3 +1365,24 @@ def test_invalid_input_data_as_object_via_main(backend: Literal["jax", "numpy"])
             orig_policy_objects={"root": METTSIM_ROOT},
             date_str="2025-01-01",
         )
+
+
+def test_request_tt_function_without_input_data(backend: Literal["jax", "numpy"]):
+    main(
+        date_str="2025-01-01",
+        main_target=MainTarget.specialized_environment.tax_transfer_function,
+        backend=backend,
+        orig_policy_objects={"root": METTSIM_ROOT},
+    )
+
+
+def test_request_raw_results_without_inputs_raises_error(
+    backend: Literal["jax", "numpy"],
+):
+    with pytest.raises(ValueError, match="The following data columns are missing"):
+        main(
+            date_str="2025-01-01",
+            main_target=MainTarget.raw_results.columns,
+            backend=backend,
+            orig_policy_objects={"root": METTSIM_ROOT},
+        )

tests/ttsim/test_end_to_end.py

@@ -1,10 +1,11 @@
 from pathlib import Path
 from typing import Literal
 
+import dags.tree as dt
 import pandas as pd
 import pytest
 
-from ttsim import InputData, MainTarget, TTTargets, main
+from ttsim import InputData, MainTarget, SpecializedEnvironment, TTTargets, main
 
 DF_WITH_NESTED_COLUMNS = pd.DataFrame(
     {
@@ -68,7 +69,7 @@
 }
 
 
-EXPECTED_RESULTS = pd.DataFrame(
+EXPECTED_TT_RESULTS = pd.DataFrame(
     {
         "payroll_tax_amount_y": [2920.0, 0.0, 0.0],
         "payroll_tax_child_tax_credit_amount_m": [8.333333, 0.0, 0.0],
@@ -99,8 +100,48 @@ def test_end_to_end(input_data_arg, backend: Literal["numpy", "jax"]):
         backend=backend,
     )
     pd.testing.assert_frame_equal(
-        EXPECTED_RESULTS,
+        EXPECTED_TT_RESULTS,
         result,
         check_dtype=False,
         check_index_type=False,
     )
+
+
+def test_can_create_input_template(backend: Literal["numpy", "jax"]):
+    result_template = main(
+        main_target=MainTarget.templates.input_data_dtypes,
+        date_str="2025-01-01",
+        orig_policy_objects={"root": Path(__file__).parent / "mettsim"},
+        backend=backend,
+        tt_targets=TTTargets(tree=TARGETS_TREE),
+    )
+    flat_result_template = dt.flatten_to_tree_paths(result_template)
+    flat_expected = dt.flatten_to_tree_paths(INPUT_DF_MAPPER)
+    assert flat_result_template.keys() == flat_expected.keys()
+
+
+def test_can_create_tt_function(backend: Literal["numpy", "jax"]):
+    tt_function = main(
+        main_target=MainTarget.specialized_environment.tax_transfer_function,
+        date_str="2025-01-01",
+        orig_policy_objects={"root": Path(__file__).parent / "mettsim"},
+        backend=backend,
+        tt_targets=TTTargets(tree=TARGETS_TREE),
+    )
+    result_df = main(
+        main_target=MainTarget.results.df_with_mapper,
+        specialized_environment=SpecializedEnvironment(
+            tax_transfer_function=tt_function,
+        ),
+        input_data=InputData.df_and_mapper(df=DF_FOR_MAPPER, mapper=INPUT_DF_MAPPER),
+        date_str="2025-01-01",
+        orig_policy_objects={"root": Path(__file__).parent / "mettsim"},
+        backend=backend,
+        tt_targets=TTTargets(tree=TARGETS_TREE),
+    )
+    pd.testing.assert_frame_equal(
+        EXPECTED_TT_RESULTS,
+        result_df,
+        check_dtype=False,
+        check_index_type=False,
+    )