adds more tests

src/sklearn_evaluation/nb/NotebookIntrospector.py

@@ -8,6 +8,11 @@
 
 
 def _safe_literal_eval(source, to_df=False, none_if_error=False):
+    """
+    Evaluates a literal, if the code cannot be parsed, it returns the original
+    source as a string unless non_if_error is True, in such case it returns
+    None
+    """
     try:
         result = ast.literal_eval(source)
 

tests/nb/conftest.py

@@ -14,6 +14,7 @@ def save_and_execute_notebook(nb_str, path):
 
     nbformat.write(nb, path)
     pm.execute_notebook(str(path), str(path))
+    return str(path)
 
 
 @pytest.fixture
@@ -89,3 +90,14 @@ def nb_no_output():
 x = 1
 """
     save_and_execute_notebook(content, 'nb_no_output.ipynb')
+
+
+@pytest.fixture
+def nb_invalid_output():
+    content = """
+import numpy as np
+
+# + tags=["numpy_array"]
+np.array([1, 2, 3])
+"""
+    return save_and_execute_notebook(content, 'nb_invalid_output.ipynb')

tests/nb/test_NotebookCollection.py

@@ -1,7 +1,7 @@
 import pytest
 
 from sklearn_evaluation import NotebookCollection
-from sklearn_evaluation.nb.NotebookCollection import HTMLMapping
+from sklearn_evaluation.nb import NotebookCollection as nbcollection
 from IPython.display import HTML
 
 
@@ -45,10 +45,64 @@ def test_custom_keys(tmp_directory, nb_literals, nb_other_literals, arg,
 
 def test_html_mapping():
     d = dict(a=1, b=2)
-    m = HTMLMapping(d, 'some html string')
+    m = nbcollection.HTMLMapping(d, 'some html string')
 
     assert list(m._ipython_key_completions_()) == ['a', 'b']
     assert list(m) == ['a', 'b']
     assert len(m) == 2
     assert m._repr_html_() == 'some html string'
     assert dict(m) == d
+
+
+def test_get_filename():
+    assert nbcollection._get_filename('/a/b/c.ipynb') == 'c'
+
+
+@pytest.mark.parametrize('elements, value, expected', [
+    [False, 'value', False],
+    [True, 'value', True],
+    [None, 'value', False],
+    [['another_value'], 'value', False],
+    [['another_value', 'value'], 'value', True],
+])
+def test_is_in(elements, value, expected):
+    assert nbcollection.is_in(elements, value) == expected
+
+
+@pytest.mark.parametrize('scores_arg, errors_expected, scores_expected', [
+    [
+        False,
+        ([1, 2, 3], ['a', 'b', 'c']),
+        ([], ['a', 'b', 'c']),
+    ],
+    [
+        True,
+        ([], ['a', 'b', 'c']),
+        ([1, 2, 3], ['a', 'b', 'c']),
+    ],
+    [
+        [1],
+        ([2, 3], ['a', 'b', 'c']),
+        ([1], ['a', 'b', 'c']),
+    ],
+])
+@pytest.mark.parametrize('transpose', [False, True])
+def test_split_errors_and_scores(scores_arg, errors_expected, scores_expected,
+                                 transpose):
+    errors, scores = nbcollection.split_errors_and_scores(
+        axis=[1, 2, 3],
+        scores_arg=scores_arg,
+        axis_second=['a', 'b', 'c'],
+        transpose=transpose)
+
+    if transpose:
+        scores_expected = (scores_expected[1], scores_expected[0])
+        errors_expected = (errors_expected[1], errors_expected[0])
+
+    assert errors == errors_expected
+    assert scores == scores_expected
+
+
+def test_data2html_img():
+    assert nbcollection.data2html_img(bytes(
+        [0])) == '<img src="data:image/png;base64, AA==\n"/>'

tests/nb/test_NotebookIntrospector.py

@@ -38,6 +38,11 @@ def test_notebook_with_no_output(tmp_directory, nb_no_output):
     assert dict(intr) == dict()
 
 
+def test_notebook_with_invalid_output(tmp_directory, nb_invalid_output):
+    intr = NotebookIntrospector(nb_invalid_output, literal_eval=True)
+    assert intr['numpy_array'] == 'array([1, 2, 3])'
+
+
 def test_json_serializable(tmp_directory, nb_plot):
     d = NotebookIntrospector('nb_plot.ipynb').to_json_serializable()