diff --git a/components/text_normalization/Dockerfile b/components/text_normalization/Dockerfile
index eb96b604a..1a3b02fd5 100644
--- a/components/text_normalization/Dockerfile
+++ b/components/text_normalization/Dockerfile
@@ -1,27 +1,30 @@
 FROM --platform=linux/amd64 python:3.8-slim as base
 
-## System dependencies
+# System dependencies
 RUN apt-get update && \
     apt-get upgrade -y && \
     apt-get install git -y
 
-RUN pip3 install fondant[aws,azure,gcp]@git+https://github.com/mrchtr/fondant
-
-# install requirements
+# Install requirements
 COPY requirements.txt /
 RUN pip3 install --no-cache-dir -r requirements.txt
 
-# Copy over src-files
+# Install Fondant
+# This is split from other requirements to leverage caching
+ARG FONDANT_VERSION=main
+RUN pip3 install fondant[aws,azure,gcp]@git+https://github.com/ml6team/fondant@${FONDANT_VERSION}
+
+# Set the working directory to the component folder
+WORKDIR /component
 COPY src/ src/
+ENV PYTHONPATH "${PYTHONPATH}:./src"
 
-# Run component tests
 FROM base as test
-RUN pip3 install pytest # TODO add pytest to package setup
-ENV PYTHONPATH "${PYTHONPATH}:./src"
+COPY test_requirements.txt .
+RUN pip3 install --no-cache-dir -r test_requirements.txt
 COPY tests/ tests/
-RUN ["python", "-m", "pytest", "tests/"]
+RUN python -m pytest tests
 
 FROM base
-# Set the working directory to the component folder
-WORKDIR /src
-ENTRYPOINT ["python", "main.py"]
\ No newline at end of file
+WORKDIR /component/src
+ENTRYPOINT ["python", "main.py"]
diff --git a/components/text_normalization/test_requirements.txt b/components/text_normalization/test_requirements.txt
new file mode 100644
index 000000000..6950eb5a7
--- /dev/null
+++ b/components/text_normalization/test_requirements.txt
@@ -0,0 +1 @@
+pytest==7.4.0
\ No newline at end of file