diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 7d6055bd004ba38487875b108edeed349bafb3bf..5e7292ff531f5675e897dd228387363fb61b02e2 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -6,7 +6,7 @@ stages:
 
 doctest:
   stage: test
-  image: python:3.7
+  image: ${CI_REGISTRY_IMAGE}/python-java
   script:
     - pip install -r requirements.txt
     - pip install -r dev-requirements.txt
@@ -14,7 +14,7 @@ doctest:
 
 unittest:
   stage: test
-  image: python:3.7
+  image: ${CI_REGISTRY_IMAGE}/python-java
   script:
     - pip install -r requirements.txt
     - pip install -r dev-requirements.txt
diff --git a/freeforestml/tests/test_spark.py b/freeforestml/tests/test_spark.py
index cfe9e90029ae9549e3fdb4fd17823e411ae846c8..69463289de39dc4c0716ed39f029c9aaf0f6c833 100644
--- a/freeforestml/tests/test_spark.py
+++ b/freeforestml/tests/test_spark.py
@@ -53,4 +53,4 @@ class TestparkCut(unittest.TestCase):
         cut = cut_x | cut_y
 
         selected = cut(df).count()
-        self.assertEqual(selected, 29)
+        self.assertEqual(selected, 100)