diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md
index 506fec4..ecdd5e0 100644
--- a/.github/PULL_REQUEST_TEMPLATE.md
+++ b/.github/PULL_REQUEST_TEMPLATE.md
@@ -23,7 +23,7 @@ Please delete options that are not relevant.
 - [ ] My code adheres to the [Google Python Style Guide](https://google.github.io/styleguide/pyguide.html)
 - [ ] I ran `make pylint` and code is rated 10/10
 - [ ] I have added tests that prove my fix is effective or that my feature works
-- [ ] I ran `make test testnb` and all tests pass
+- [ ] I ran `make test` and all tests pass
 - [ ] I ran the tool and verified the change works
 - [ ] I have adequately commented my code, particularly in hard-to-understand areas
 - [ ] I have made relevant changes to the documentation, if needed
diff --git a/.github/workflows/python-cicd.yml b/.github/workflows/python-cicd.yml
index 3fefdae..9d80335 100644
--- a/.github/workflows/python-cicd.yml
+++ b/.github/workflows/python-cicd.yml
@@ -23,13 +23,9 @@ jobs:
         python -m pip install --upgrade pip
         if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
 
-    - name: Run Jupyter notebook tests
+    - name: Run all tests
       run: |
         export PYTHONPATH="$GITHUB_WORKSPACE"
-        make testnb
-
-    - name: Run Python tests
-      run: |
         make test
 
     - name: Lint with pylint
diff --git a/Makefile b/Makefile
index 7024d82..f69cb7e 100644
--- a/Makefile
+++ b/Makefile
@@ -3,13 +3,15 @@ all: init testnb test pylint
 init:
 	pip install -r requirements.txt
 
-test:
+test: test-nb test-py
+
+test-py:
 	nosetests --with-coverage -v --cover-package=tfrecorder
 
-testnb:
+test-nb:
 	ls -1 samples/*.ipynb | grep -v '^.*Dataflow.ipynb' | xargs py.test --nbval-lax -p no:python
 
 pylint:
 	pylint -j 0 tfrecorder
 
-.PHONY: all init testnb test pylint
+.PHONY: all init test pylint
diff --git a/tfrecorder/utils_test.py b/tfrecorder/utils_test.py
index 15dd51a..77331e4 100644
--- a/tfrecorder/utils_test.py
+++ b/tfrecorder/utils_test.py
@@ -31,7 +31,7 @@
 from tfrecorder import utils
 from tfrecorder import test_utils
 from tfrecorder import input_schema
-from tfrecorder import dataset_loader as _dataset
+from tfrecorder import dataset_loader
 
 
 # pylint: disable=protected-access
@@ -69,7 +69,7 @@ def setUp(self):
     self.data = data
     self.dataset = tf.data.Dataset.from_tensor_slices(self.data)
 
-  @mock.patch.object(_dataset, 'load', autospec=True)
+  @mock.patch.object(dataset_loader, 'load', autospec=True)
   def test_valid_records(self, mock_fn):
     """Tests valid case on reading multiple records."""
 
@@ -98,7 +98,7 @@ def test_valid_records(self, mock_fn):
       expected_image_files = self.data['image_name']
       self.assertCountEqual(actual_image_files, expected_image_files)
 
-  @mock.patch.object(_dataset, 'load', autospec=True)
+  @mock.patch.object(dataset_loader, 'load', autospec=True)
   def test_no_data_for_split(self, mock_fn):
     """Check exception raised when data could not be loaded given `split`."""