From 6b175fd518b0d8318dd2c08e7a19c8ac714f13a5 Mon Sep 17 00:00:00 2001
From: Danilo Lessa Bernardineli <danilo.lessa@gmail.com>
Date: Thu, 18 Apr 2024 16:09:51 -0300
Subject: [PATCH] conditionally disable tqdm + add print test

---
 cadCAD/engine/__init__.py |  6 ++-
 testing/test_print.py     | 77 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 2 deletions(-)
 create mode 100644 testing/test_print.py

diff --git a/cadCAD/engine/__init__.py b/cadCAD/engine/__init__.py
index 5e08856..8ce480a 100644
--- a/cadCAD/engine/__init__.py
+++ b/cadCAD/engine/__init__.py
@@ -104,7 +104,8 @@ def execute(self) -> Tuple[object, object, Dict[str, object]]:
         t1 = time()
         for x in tqdm(self.configs,
                       total=len(self.configs),
-                      desc="Initializing configurations"):
+                      desc="Initializing configurations",
+                      disable=self.supress_print):
             sessions.append(
                 {
                     'user_id': x.user_id, 'experiment_id': x.experiment_id, 'session_id': x.session_id,
@@ -182,7 +183,8 @@ def get_final_results(simulations: List[StateHistory],
             flat_timesteps, tensor_fields = [], []
             for sim_result, psu, ep in tqdm(list(zip(simulations, psus, eps)),
                                             total=len(simulations),
-                                            desc='Flattening results'):
+                                            desc='Flattening results',
+                                            disable=self.supress_print):
                 if do_flatten:
                     flat_timesteps.append(flatten(sim_result))
                 tensor_fields.append(create_tensor_field(psu, ep))
diff --git a/testing/test_print.py b/testing/test_print.py
new file mode 100644
index 0000000..07c0e62
--- /dev/null
+++ b/testing/test_print.py
@@ -0,0 +1,77 @@
+from cadCAD.configuration import Experiment
+from cadCAD.configuration.utils import config_sim
+from cadCAD.engine import Executor, ExecutionContext, ExecutionMode
+import pytest
+
+P_no_lst = {'pA': 1, 'pB': 2, 'pC': 3}
+P_single_lst = {'pA': [1], 'pB': [1], 'pC': [3]}
+P_single_swp = {'pA': [4, 5, 6], 'pB': [1], 'pC': [3]}
+P_all_swp = {'pA': [7, 8, 9], 'pB': [1, 2, 3], 'pC': [1, 2, 3]}
+P_all_but_one_swp = {'pA': [7, 8, 9], 'pB': [1, 2, 3], 'pC': [1]}
+Ps = [P_no_lst, P_single_lst, P_single_swp, P_all_swp, P_all_but_one_swp]
+
+CONFIG_SIGNATURES_TO_TEST = [(3, 3, 3, 3, 3),
+                             (1, 3, 3, 3, 3),
+                             (3, 1, 3, 3, 3),
+                             (1, 1, 3, 3, 3),
+                             (3, 3, 1, 3, 3),
+                             (1, 3, 1, 3, 3),
+                             (1, 1, 1, 3, 3)]
+
+
+def run_experiment(exp: Experiment, mode: str, supress_print=False):
+    exec_context = ExecutionContext(mode)
+    executor = Executor(exec_context=exec_context, configs=exp.configs, supress_print=supress_print)
+    (records, tensor_field, _) = executor.execute()
+    return records
+
+
+def param_count_test_suf_generator(provided_params):
+    def s_test_param_count(params, _2, _3, _4, _5):
+        assert params.keys() == provided_params.keys(), 'Params are not matching'
+        return ('varA', None)
+    return s_test_param_count
+
+
+def param_count_test_policy_generator(provided_params):
+    def p_test_param_count(params, _2, _3, _4):
+        assert params.keys() == provided_params.keys(), 'Params are not matching'
+        return {'sigA': None}
+    return p_test_param_count
+
+
+def create_experiments(N_simulations=3, N_sweeps=3, N_runs=3, N_timesteps=3, N_substeps=3, params={}) -> Experiment:
+
+    INITIAL_STATE = {'varA': None}
+    PSUBs = [{'policies': {'sigA': param_count_test_policy_generator(
+        params)}, 'variables': {'varA': param_count_test_suf_generator(params)}}] * N_substeps
+
+    SIM_CONFIG = config_sim(
+        {
+            "N": N_runs,
+            "T": range(N_timesteps),
+            "M": params,  # Optional
+        }
+    )
+
+    exp = Experiment()
+    for i_sim in range(N_simulations):
+        exp.append_model(
+            sim_configs=SIM_CONFIG,
+            initial_state=INITIAL_STATE,
+            partial_state_update_blocks=PSUBs
+        )
+    return exp
+
+
+
+def test_print(capfd):
+    exp = run_experiment(create_experiments(N_simulations=3, N_sweeps=3, N_runs=3, N_timesteps=3, N_substeps=3, params={'a': 0}), 'single_proc', supress_print=False)
+    out, err = capfd.readouterr()
+    assert "                  ___________    ____\n  ________ __ ___/ / ____/   |  / __ \\\n / ___/ __` / __  / /   / /| | / / / /\n/ /__/ /_/ / /_/ / /___/ ___ |/ /_/ /\n\\___/\\__,_/\\__,_/\\____/_/  |_/_____/\nby cadCAD" in out
+    assert 'Initializing configurations' in err
+
+    exp = run_experiment(create_experiments(N_simulations=3, N_sweeps=3, N_runs=3, N_timesteps=3, N_substeps=3, params={'a': 0}), 'single_proc', supress_print=True)
+    out, err = capfd.readouterr()
+    assert out == ''
+    assert err == ''
\ No newline at end of file