From c7298fad444ac758ebfa04c7153803656c1fb909 Mon Sep 17 00:00:00 2001
From: Jeff Omhover <jf.omhover@gmail.com>
Date: Thu, 9 Dec 2021 10:20:14 -0800
Subject: [PATCH] Extend coverage: add unit tests where they miss most (#195)

* add unit tests for lightgbm utils
* add coverage config, remove tasks from coverage
* add tests for data2bin
* remove distributed tests for now
---
 .coveragerc                             |  5 +++
 tests/common/test_lightgbm_utils.py     | 48 +++++++++++++++++++++++++
 tests/scripts/test_lightgbm_data2bin.py | 37 +++++++++++++++++++
 3 files changed, 90 insertions(+)
 create mode 100644 .coveragerc
 create mode 100644 tests/common/test_lightgbm_utils.py
 create mode 100644 tests/scripts/test_lightgbm_data2bin.py

diff --git a/.coveragerc b/.coveragerc
new file mode 100644
index 00000000..908b4d7e
--- /dev/null
+++ b/.coveragerc
@@ -0,0 +1,5 @@
+[report]
+exclude_lines =
+    pragma: no cover
+omit =
+    src/common/tasks.py
diff --git a/tests/common/test_lightgbm_utils.py b/tests/common/test_lightgbm_utils.py
new file mode 100644
index 00000000..21f10b19
--- /dev/null
+++ b/tests/common/test_lightgbm_utils.py
@@ -0,0 +1,48 @@
+"""Tests src/common/io.py"""
+import os
+import pytest
+from unittest.mock import call, Mock, patch
+
+from common.lightgbm_utils import LightGBMCallbackHandler
+from lightgbm.callback import CallbackEnv
+
+def test_lightgbm_callback_handler():
+    metrics_logger = Mock()
+
+    callback_handler = LightGBMCallbackHandler(
+        metrics_logger, metrics_prefix=None, metrics_suffix=None
+    )
+
+    # namedtuple
+    # see https://lightgbm.readthedocs.io/en/latest/_modules/lightgbm/callback.html
+    callback_env = CallbackEnv(
+        None, # model
+        {"foo_param": 0.32}, # params
+        3, # iteration
+        0, # begin_iteration
+        5, # end_iteration
+        [
+            # list of tuples
+            (
+                "valid_0", # dataset name
+                "rmse", # evaluation name
+                12345.0, # result
+                None, # _
+            ),
+            (
+                "valid_0", # dataset name
+                "l2", # evaluation name
+                3456.0, # result
+                None, # _
+            )
+        ]
+    )
+    callback_handler.callback(callback_env)
+
+    metrics_logger.log_metric.assert_has_calls(
+        [
+            call(key="valid_0.rmse", value=12345.0, step=3),
+            call(key="valid_0.l2", value=3456.0, step=3)
+        ],
+        any_order=True
+    )
diff --git a/tests/scripts/test_lightgbm_data2bin.py b/tests/scripts/test_lightgbm_data2bin.py
new file mode 100644
index 00000000..cd471dcc
--- /dev/null
+++ b/tests/scripts/test_lightgbm_data2bin.py
@@ -0,0 +1,37 @@
+"""
+test src/scripts/partition_data/partition.py
+"""
+import os
+import sys
+import tempfile
+import pytest
+from unittest.mock import patch
+
+from scripts.data_processing.lightgbm_data2bin import data2bin
+
+# IMPORTANT: see conftest.py for fixtures
+
+def test_lightgbm_data2bin(temporary_dir, regression_train_sample, regression_test_sample):
+    """Tests src/scripts/data_processing/lightgbm_data2bin/data2bin.py"""
+    binary_train_data_dir = os.path.join(temporary_dir, "binary_train_data")
+    binary_test_data_dir = os.path.join(temporary_dir, "binary_test_data")
+
+    # create test arguments for the script
+
+    script_args = [
+        "data2bin.py",
+        "--train", regression_train_sample,
+        "--test", regression_test_sample,
+        "--output_train", binary_train_data_dir,
+        "--output_test", binary_test_data_dir,
+        "--header", "False",
+        "--label_column", "0",
+        "--max_bin", "255",
+    ]
+
+    # replaces sys.argv with test arguments and run main
+    with patch.object(sys, "argv", script_args):
+        data2bin.main()
+
+    assert os.path.isfile(os.path.join(binary_train_data_dir, "train.bin"))
+    assert os.path.isfile(os.path.join(binary_test_data_dir, "test_0.bin"))