diff --git a/.github/workflows/self_test.yaml b/.github/workflows/self_test.yaml
index 66a95dc..16ae6fe 100644
--- a/.github/workflows/self_test.yaml
+++ b/.github/workflows/self_test.yaml
@@ -22,13 +22,3 @@ jobs:
           app_version_path: "version.txt"
           docker_compose_path: "docker-compose.yml"
           labels: ${{ toJSON(github.event.pull_request.labels.*.name) }}
-
-      - name: Log Success
-        if: ${{ env.app_updated == 'true' }}
-        run: |
-          echo "App version has been updated correctly!"
-
-      - name: Log Success
-        if: ${{ env.compose_updated == 'true' }}
-        run: |
-          echo "Compose version has been updated correctly!"
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
index 904c4c4..890fd30 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,7 @@
 .idea
-main_version
-branch_version
 src/__pycache__/*
-tests/__pycache__/*
\ No newline at end of file
+src/features/__pycache__/*
+tests/__pycache__/*
+.coverage
+.coveragerc
+coverage.xml
\ No newline at end of file
diff --git a/.pylintrc b/.pylintrc
index fad2aa6..fa8b34b 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -5,9 +5,9 @@ max-line-length=120
 
 # Disable various warnings:
 
-# W0237 Disabled as it makes the code more readable
-# R0801 Disabled as it's a small amount of duplication
-disable=W0237, R0801
+# Disable R0801 as the code is only repeated twice.
+
+disable=R0801
 
 [MASTER]
 init-hook='import sys; sys.path.append("src")'
\ No newline at end of file
diff --git a/README.md b/README.md
index a710f39..bea978a 100644
--- a/README.md
+++ b/README.md
@@ -5,9 +5,9 @@
 [![codecov](https://codecov.io/gh/stfc/check-version-action/graph/badge.svg?token=OD2Z90ST8R)](https://codecov.io/gh/stfc/check-version-action)
 
 
-This action compares the application version number from your working branch to the main branch.
+This action compares the app version number from your working branch to the main branch.
 
-You can also check that the **first** image version that appears in your `docker-compose.yaml` file will match the application version
+You can also check that the **first** image version that appears in your `docker-compose.yaml` file matches the app version
 
 The comparison follows the PEP 440 Version Identification and Dependency Specification.
 
@@ -42,27 +42,16 @@ If you are making a change which should not affect the version such as README or
     path: 'branch'
     
 - name: Compare versions
-  # Don't run on main otherwise it will compare main with main
+  # Don't run on main otherwise it compares main with main
   if: ${{ github.ref != 'refs/heads/main' }} 
   id: version_comparison
   uses: stfc/check-version-action@main
   with:
     # Path to version file from project root
     app_version_path: "version.txt"
-    # Optional: To check if compose image version matches application version
+    # Optional: to check if Docker compose image version matches app version
     docker_compose_path: "docker-compose.yaml"
     labels: ${{ toJSON(github.event.pull_request.labels.*.name) }}
-    
-- name: Log App Success
-  if: ${{ env.app_updated == 'true' }}
-  run: |
-    echo "App version has been updated correctly!"
-
-# Optional: If using the docker compose check
-- name: Log Compose Success
-  if: ${{ env.compose_updated == 'true' }}
-  run: |
-    echo "Compose version has been updated correctly!"
 ```
 <!-- end usage -->
 
diff --git a/docker-compose.yml b/docker-compose.yml
index b8fb975..86d4c8f 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -1,3 +1,3 @@
 services:
   self-test:
-    image: some/test:1.1.0
+    image: some/test:1.2.0
diff --git a/src/base.py b/src/base.py
deleted file mode 100644
index 84a79f1..0000000
--- a/src/base.py
+++ /dev/null
@@ -1,36 +0,0 @@
-"""This is the base module including helper/abstract classes and errors."""
-
-from abc import ABC, abstractmethod
-from pathlib import Path
-from typing import Union
-from packaging.version import Version
-
-
-class VersionNotUpdated(Exception):
-    """The version number has not been updated or updated incorrectly."""
-
-
-class Base(ABC):
-    """The base abstract class to build features on."""
-
-    @abstractmethod
-    def run(self, path1: Path, path2: Path) -> bool:
-        """
-        This method is the entry point to the feature.
-        It should take two paths and return the comparison result.
-        """
-
-    @staticmethod
-    @abstractmethod
-    def read_files(path1: Path, path2: Path) -> (str, str):
-        """This method should read the contents of the compared files and return the strings"""
-
-    @staticmethod
-    @abstractmethod
-    def get_version(content: str) -> Version:
-        """This method should extract the version from the file and return it as a packaging Version object"""
-
-    @staticmethod
-    @abstractmethod
-    def compare(version1: Version, version2: Version) -> Union[bool, VersionNotUpdated]:
-        """This method should compare the versions and return a bool status"""
diff --git a/src/comparison.py b/src/comparison.py
deleted file mode 100644
index 6cc1c3c..0000000
--- a/src/comparison.py
+++ /dev/null
@@ -1,127 +0,0 @@
-"""The comparison module which is where the main check classes are."""
-
-from pathlib import Path
-from typing import Union, List, Type
-
-from packaging.version import Version
-from base import Base, VersionNotUpdated
-
-
-class CompareAppVersion(Base):
-    """This class compares the application versions"""
-
-    def run(self, path1: Path, path2: Path) -> bool:
-        """
-        Entry point to compare application versions.
-        :param path1: Path to main version
-        :param path2: Path to branch version
-        :return: true if success, error if fail
-        """
-        main_content, branch_content = self.read_files(path1, path2)
-        main_ver = self.get_version(main_content)
-        branch_ver = self.get_version(branch_content)
-        comparison = self.compare(main_ver, branch_ver)
-        if comparison == VersionNotUpdated:
-            raise VersionNotUpdated(
-                f"The version in {('/'.join(str(path2).split('/')[4:]))[0:]} has not been updated correctly."
-            )
-        return True
-
-    @staticmethod
-    def read_files(path1: Path, path2: Path) -> (str, str):
-        """
-        Read both version files and return the contents
-        :param path1: Path to main version
-        :param path2: Path to branched version
-        :return: main_ver, branch_ver
-        """
-        with open(path1, "r", encoding="utf-8") as file1:
-            content1 = file1.read()
-        with open(path2, "r", encoding="utf-8") as file2:
-            content2 = file2.read()
-        return content1, content2
-
-    @staticmethod
-    def get_version(content: str) -> Version:
-        """
-        This method returns the version from the file as an object
-        For application versions we expect nothing else in the file than the version.
-        :param content: Application version string
-        :return: Application version object
-        """
-        return Version(content)
-
-    @staticmethod
-    def compare(main: Version, branch: Version) -> Union[bool, Type[VersionNotUpdated]]:
-        """
-        Returns if the branch version is larger than the main version
-        :param main: Version on main
-        :param branch: Version on branch
-        :return: If the version update is correct return true, else return error
-        """
-        if branch > main:
-            return True
-        return VersionNotUpdated
-
-
-class CompareComposeVersion(Base):
-    """This class compares the docker compose image version to the application version."""
-
-    def run(self, app: Path, compose: Path) -> bool:
-        """
-        Entry point to compare docker compose and application versions.
-        :param app: Path to application version
-        :param compose: Path to compose image version
-        :return: true if success, error if fail
-        """
-        app_content, compose_content = self.read_files(app, compose)
-        app_ver = Version(app_content)
-        compose_ver = self.get_version(compose_content)
-        comparison = self.compare(app_ver, compose_ver)
-        if comparison == VersionNotUpdated:
-            raise VersionNotUpdated(
-                f"The version in {('/'.join(str(compose).split('/')[4:]))[0:]}"
-                f"does not match {('/'.join(str(app).split('/')[4:]))[0:]}."
-            )
-        return True
-
-    @staticmethod
-    def read_files(app: Path, compose: Path) -> (str, List):
-        """
-        Read both version files and return the contents
-        :param app: Path to app version
-        :param compose: Path to compose version
-        :return: main_ver, branch_ver
-        """
-        with open(app, "r", encoding="utf-8") as file1:
-            content1 = file1.read()
-        with open(compose, "r", encoding="utf-8") as file2:
-            content2 = file2.readlines()
-        return content1, content2
-
-    @staticmethod
-    def get_version(content: List[str]) -> Version:
-        """
-        This method returns the version from the file as an object
-        For compose versions we have to do some data handling.
-        :param content: Compose version string
-        :return: Compose version object
-        """
-        version_str = ""
-        for line in content:
-            if "image" in line:
-                version_str = line.strip("\n").split(":")[-1]
-                break
-        return Version(version_str)
-
-    @staticmethod
-    def compare(app: Version, compose: Version) -> Union[bool, Type[VersionNotUpdated]]:
-        """
-        Returns if the application version and docker compose version are equal.
-        :param app: App version
-        :param compose: Compose version
-        :return: If the version update is correct return true, else return error
-        """
-        if app == compose:
-            return True
-        return VersionNotUpdated
diff --git a/src/features/__init__.py b/src/features/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/features/app_version.py b/src/features/app_version.py
new file mode 100644
index 0000000..322c69d
--- /dev/null
+++ b/src/features/app_version.py
@@ -0,0 +1,60 @@
+"""Compare app version.txt on main to the branch."""
+
+from pathlib import Path
+
+from packaging.version import Version
+
+
+class CompareAppVersion:
+    """This class compares the app versions"""
+
+    def run(self, path1: Path, path2: Path) -> bool:
+        """
+        Entry point to compare app versions.
+        :param path1: Path to main version
+        :param path2: Path to branch version
+        :return: true if success, error if fail
+        """
+        main_content, branch_content = self.read_files(path1, path2)
+        main_ver = self.get_version(main_content)
+        branch_ver = self.get_version(branch_content)
+        comparison = self.compare(main_ver, branch_ver)
+        if not comparison:
+            raise RuntimeError(
+                f"The version in {('/'.join(str(path2).split('/')[4:]))[0:]} has not been updated correctly."
+            )
+        return True
+
+    @staticmethod
+    def read_files(path1: Path, path2: Path) -> (str, str):
+        """
+        Read both version files and return the contents
+        :param path1: Path to main version
+        :param path2: Path to branched version
+        :return: main_ver, branch_ver
+        """
+        with open(path1, "r", encoding="utf-8") as file1:
+            content1 = file1.read()
+        with open(path2, "r", encoding="utf-8") as file2:
+            content2 = file2.read()
+        return content1, content2
+
+    @staticmethod
+    def get_version(content: str) -> Version:
+        """
+        This method returns the version from the file as an object
+        For app versions we expect nothing else in the file than the version.
+        :param content: app version string
+        :return: app version object
+        """
+        return Version(content)
+
+    @staticmethod
+    def compare(main: Version, branch: Version) -> bool:
+        """
+        Returns if the branch version is larger than the main version
+        :param main: Version on main
+        :param branch: Version on branch
+        :return: If the version update is correct return true, else return error
+        """
+        return branch > main
diff --git a/src/features/compose_version.py b/src/features/compose_version.py
new file mode 100644
index 0000000..66ad62a
--- /dev/null
+++ b/src/features/compose_version.py
@@ -0,0 +1,67 @@
+"""Compare Docker compose image version to the version.txt."""
+
+from pathlib import Path
+from typing import List
+
+from packaging.version import Version
+
+
+class CompareComposeVersion:
+    """This class compares the docker compose image version to the app version."""
+
+    def run(self, app: Path, compose: Path) -> bool:
+        """
+        Entry point to compare docker compose and app versions.
+        :param app: Path to app version
+        :param compose: Path to compose image version
+        :return: true if success, error if fail
+        """
+        app_content, compose_content = self.read_files(app, compose)
+        app_ver = Version(app_content)
+        compose_ver = self.get_version(compose_content)
+        comparison = self.compare(app_ver, compose_ver)
+        if not comparison:
+            raise RuntimeError(
+                f"The version in {('/'.join(str(compose).split('/')[4:]))[0:]}"
+                f" does not match {('/'.join(str(app).split('/')[4:]))[0:]}."
+            )
+        return True
+
+    @staticmethod
+    def read_files(app: Path, compose: Path) -> (str, List):
+        """
+        Read both version files and return the contents
+        :param app: Path to app version
+        :param compose: Path to compose version
+        :return: main_ver, branch_ver
+        """
+        with open(app, "r", encoding="utf-8") as file1:
+            content1 = file1.read()
+        with open(compose, "r", encoding="utf-8") as file2:
+            content2 = file2.readlines()
+        return content1, content2
+
+    @staticmethod
+    def get_version(content: List[str]) -> Version:
+        """
+        This method returns the version from the file as an object
+        For compose versions we have to do some data handling.
+        :param content: Compose version string
+        :return: Compose version object
+        """
+        version_str = ""
+        for line in content:
+            if "image" in line:
+                version_str = line.strip("\n").split(":")[-1]
+                break
+        return Version(version_str)
+
+    @staticmethod
+    def compare(app: Version, compose: Version) -> bool:
+        """
+        Returns if the app version and docker compose version are equal.
+        :param app: App version
+        :param compose: Compose version
+        :return: If the version update is correct return true, else return error
+        """
+        return app == compose
diff --git a/src/main.py b/src/main.py
index 34cf5f4..543de3c 100644
--- a/src/main.py
+++ b/src/main.py
@@ -2,7 +2,8 @@
 
 import os
 from pathlib import Path
-from comparison import CompareAppVersion, CompareComposeVersion
+from src.features.compose_version import CompareComposeVersion
+from src.features.app_version import CompareAppVersion
 
 
 def main() -> bool:
@@ -10,30 +11,32 @@ def main() -> bool:
     The entry point function for the action.
     Here we get environment variables then set environment variables when finished.
     """
+    # Check if the action should skip version checks
+    for label in os.environ.get("INPUT_LABELS"):
+        if label in ["documentation", "workflow"]:
+            return False
+
+    # Collect various paths from the environment
     app_path = Path(os.environ.get("INPUT_APP_VERSION_PATH"))
-    compose_path = os.environ.get("INPUT_DOCKER_COMPOSE_PATH")
+    compose_path = Path(os.environ.get("INPUT_DOCKER_COMPOSE_PATH"))
     root_path = Path(os.environ.get("GITHUB_WORKSPACE"))
     main_path = root_path / "main"
     branch_path = root_path / "branch"
-    with open(branch_path / app_path, "r", encoding="utf-8") as release_file:
-        release_version = release_file.read().strip("\n")
-
-    labels = os.environ.get("INPUT_LABELS")
-    if any(label in labels for label in ["documentation", "workflow"]):
-        return False
 
+    # Action must compare the app version as the minimum feature.
     CompareAppVersion().run(main_path / app_path, branch_path / app_path)
+
+    # Compare the Docker compose file version if given
     if compose_path:
-        compose_path = Path(compose_path)
         CompareComposeVersion().run(branch_path / app_path, branch_path / compose_path)
 
-    github_env = os.getenv("GITHUB_ENV")
-    with open(github_env, "a", encoding="utf-8") as env:
-        # We can assume either/both of these values returned true otherwise they would have errored
+    with open(os.getenv("GITHUB_ENV"), "a", encoding="utf-8") as env:
+        # We can assume these values returned true otherwise they would have raised an error.
         env.write("app_updated=true\n")
+        # If Docker compose path was provided we can assume it returned true otherwise there would be an error.
         if compose_path:
             env.write("compose_updated=true")
-        env.write(f"release_tag={release_version}")
+
     return True
 
 
diff --git a/tests/test_compare_app_version.py b/tests/test_app_version.py
similarity index 69%
rename from tests/test_compare_app_version.py
rename to tests/test_app_version.py
index 02d8459..fb15d3d 100644
--- a/tests/test_compare_app_version.py
+++ b/tests/test_app_version.py
@@ -1,10 +1,10 @@
-"""Tests for comparison.CompareAppVersion"""
+"""Tests for features.app_version.CompareAppVersion"""
 
 from unittest.mock import patch, mock_open
 from pathlib import Path
 from packaging.version import Version
 import pytest
-from comparison import CompareAppVersion, VersionNotUpdated
+from features.app_version import CompareAppVersion
 
 
 @pytest.fixture(name="instance", scope="function")
@@ -13,9 +13,9 @@ def instance_fixture():
     return CompareAppVersion()
 
 
-@patch("comparison.CompareAppVersion.compare")
-@patch("comparison.CompareAppVersion.get_version")
-@patch("comparison.CompareAppVersion.read_files")
+@patch("features.app_version.CompareAppVersion.compare")
+@patch("features.app_version.CompareAppVersion.get_version")
+@patch("features.app_version.CompareAppVersion.read_files")
 def test_run(mock_read, mock_get_version, mock_compare, instance):
     """Test the run method makes correct calls."""
     mock_path1 = Path("mock1")
@@ -31,14 +31,14 @@ def test_run(mock_read, mock_get_version, mock_compare, instance):
     assert res
 
 
-@patch("comparison.CompareAppVersion.compare")
-@patch("comparison.CompareAppVersion.get_version")
-@patch("comparison.CompareAppVersion.read_files")
+@patch("features.app_version.CompareAppVersion.compare")
+@patch("features.app_version.CompareAppVersion.get_version")
+@patch("features.app_version.CompareAppVersion.read_files")
 def test_run_fails(mock_read, _, mock_compare, instance):
     """Test the run method fails."""
     mock_read.return_value = ("mock1", "mock2")
-    mock_compare.side_effect = VersionNotUpdated()
-    with pytest.raises(VersionNotUpdated):
+    mock_compare.side_effect = RuntimeError()
+    with pytest.raises(RuntimeError):
         instance.run(Path("mock1"), Path("mock2"))
 
 
@@ -56,12 +56,12 @@ def test_get_version(instance):
 
 
 def test_compare_pass(instance):
-    """Test that the compare returns true for a valid comparison"""
+    """Test that the compare returns true for a valid features.app_version"""
     res = instance.compare(Version("1.0.0"), Version("1.0.1"))
-    assert res != VersionNotUpdated
+    assert res
 
 
 def test_compare_fails(instance):
-    """Test that the compare returns an error for an invalid comparison"""
+    """Test that the compare returns an error for an invalid features.app_version"""
     res = instance.compare(Version("1.0.1"), Version("1.0.0"))
-    assert res == VersionNotUpdated
+    assert not res
diff --git a/tests/test_compare_compose_version.py b/tests/test_compose_version.py
similarity index 67%
rename from tests/test_compare_compose_version.py
rename to tests/test_compose_version.py
index 2874506..c3ea785 100644
--- a/tests/test_compare_compose_version.py
+++ b/tests/test_compose_version.py
@@ -1,10 +1,10 @@
-"""Tests for comparison.CompareComposeVersion"""
+"""Tests for features/compose_version.py"""
 
 from unittest.mock import patch, mock_open
 from pathlib import Path
 import pytest
 from packaging.version import Version
-from comparison import CompareComposeVersion, VersionNotUpdated
+from features.compose_version import CompareComposeVersion
 
 
 @pytest.fixture(name="instance", scope="function")
@@ -13,9 +13,9 @@ def instance_fixture():
     return CompareComposeVersion()
 
 
-@patch("comparison.CompareComposeVersion.compare")
-@patch("comparison.CompareComposeVersion.get_version")
-@patch("comparison.CompareComposeVersion.read_files")
+@patch("features.compose_version.CompareComposeVersion.compare")
+@patch("features.compose_version.CompareComposeVersion.get_version")
+@patch("features.compose_version.CompareComposeVersion.read_files")
 def test_run(mock_read, mock_get_version, mock_compare, instance):
     """Test the run method makes correct calls."""
     mock_path1 = Path("mock1")
@@ -30,14 +30,14 @@ def test_run(mock_read, mock_get_version, mock_compare, instance):
     assert res
 
 
-@patch("comparison.CompareComposeVersion.compare")
-@patch("comparison.CompareComposeVersion.get_version")
-@patch("comparison.CompareComposeVersion.read_files")
+@patch("features.compose_version.CompareComposeVersion.compare")
+@patch("features.compose_version.CompareComposeVersion.get_version")
+@patch("features.compose_version.CompareComposeVersion.read_files")
 def test_run_fails(mock_read, _, mock_compare, instance):
     """Test the run method fails."""
     mock_read.return_value = ("1.0.0", "1.0.1")
-    mock_compare.side_effect = VersionNotUpdated()
-    with pytest.raises(VersionNotUpdated):
+    mock_compare.side_effect = RuntimeError()
+    with pytest.raises(RuntimeError):
         instance.run(Path("mock1"), Path("mock2"))
 
 
@@ -55,12 +55,12 @@ def test_get_version(instance):
 
 
 def test_compare_pass(instance):
-    """Test that the compare returns true for a valid comparison"""
+    """Test that the compare returns true for a valid features.compose_version"""
     res = instance.compare(Version("1.0.0"), Version("1.0.0"))
-    assert res != VersionNotUpdated
+    assert res
 
 
 def test_compare_fails(instance):
-    """Test that the compare returns an error for an invalid comparison"""
+    """Test that the compare returns an error for an invalid features.compose_version"""
     res = instance.compare(Version("1.0.1"), Version("1.0.0"))
-    assert res == VersionNotUpdated
+    assert not res
diff --git a/tests/test_main.py b/tests/test_main.py
index 810f25a..6fbf333 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -1,4 +1,4 @@
-"""Tests for main"""
+"""Tests for main.py"""
 
 from unittest.mock import patch, mock_open
 from pathlib import Path
@@ -11,10 +11,10 @@
 def test_main(mock_os, mock_compare_app, mock_compare_compose):
     """Test the main method runs correctly."""
     mock_os.environ.get.side_effect = [
+        ["some_label"],
         Path("app"),
         Path("compose"),
         Path("workspace"),
-        [],
     ]
 
     with patch("builtins.open", mock_open(read_data="1.0.0")):
@@ -38,10 +38,10 @@ def test_main(mock_os, mock_compare_app, mock_compare_compose):
 def test_main_skip(mock_os):
     """Test the main method skips the comparison methods."""
     mock_os.environ.get.side_effect = [
+        ["workflow", "documentation"],
         Path("app"),
         Path("compose"),
         Path("workspace"),
-        ["workflow", "documentation"],
     ]
     with patch("builtins.open", mock_open(read_data="1.0.0")):
         res = main()
diff --git a/version.txt b/version.txt
index 9084fa2..26aaba0 100644
--- a/version.txt
+++ b/version.txt
@@ -1 +1 @@
-1.1.0
+1.2.0