diff --git a/pain001/__main__.py b/pain001/__main__.py
index 0bc5fcb..005703b 100644
--- a/pain001/__main__.py
+++ b/pain001/__main__.py
@@ -1,38 +1,10 @@
-# Copyright (C) 2023 Sebastien Rousseau.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
-# implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# pylint: disable=invalid-name
-"""
-Enables use of Python Pain001 as a "main" function (i.e.
-"python3 -m pain001
-<xml_message_type> <xml_template_file_path> <xsd_schema_file_path>
-<data_file_path>").
-
-This allows using Pain001 with third-party libraries without modifying
-their code.
-"""
-
-
+# Other imports remain the same
+import click
 import os
 import sys
-import click
-
 from pain001.constants.constants import valid_xml_types
 from pain001.context.context import Context
 from pain001.core.core import process_files
-
 from rich.console import Console
 from rich.table import Table
 from rich import box
@@ -47,7 +19,6 @@
 title = "Pain001"
 
 table = Table(box=box.ROUNDED, safe_box=True, show_header=False, title=title)
-
 table.add_column(justify="center", no_wrap=False, vertical="middle")
 table.add_row(description)
 table.width = 80
@@ -85,6 +56,20 @@
     type=click.Path(),
     help="Path to data file (CSV or SQLite) (required)",
 )
+def cli(
+    xml_message_type,
+    xml_template_file_path,
+    xsd_schema_file_path,
+    data_file_path,
+):
+    main(
+        xml_message_type,
+        xml_template_file_path,
+        xsd_schema_file_path,
+        data_file_path,
+    )
+
+
 def main(
     xml_message_type,
     xml_template_file_path,
@@ -163,4 +148,4 @@ def main(
 
 if __name__ == "__main__":
     # pylint: disable=no-value-for-parameter
-    main()
+    cli()
diff --git a/tests/test_main.py b/tests/test_main.py
index 49ea99f..c15402c 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -1,5 +1,5 @@
 from click.testing import CliRunner
-from pain001.__main__ import main
+from pain001.__main__ import cli
 
 
 class TestMain:
@@ -12,7 +12,7 @@ def setup_method(self):
 
     def test_main_with_valid_files(self):
         result = self.runner.invoke(
-            main,
+            cli,
             [
                 "--xml_message_type",
                 self.xml_message_type,
@@ -32,7 +32,7 @@ def test_main_with_valid_files(self):
 
     def test_main_with_missing_xml_message_type(self):
         result = self.runner.invoke(
-            main,
+            cli,
             [
                 "--xml_template_file_path",
                 self.xml_file,
@@ -47,7 +47,7 @@ def test_main_with_missing_xml_message_type(self):
 
     def test_main_with_missing_xsd_template_file(self):
         result = self.runner.invoke(
-            main,
+            cli,
             [
                 "--xml_message_type",
                 self.xml_message_type,
@@ -62,7 +62,7 @@ def test_main_with_missing_xsd_template_file(self):
 
     def test_main_with_missing_data_file(self):
         result = self.runner.invoke(
-            main,
+            cli,
             [
                 "--xml_message_type",
                 self.xml_message_type,
@@ -77,7 +77,7 @@ def test_main_with_missing_data_file(self):
 
     def test_main_with_invalid_xml_message_type(self):
         result = self.runner.invoke(
-            main,
+            cli,
             [
                 "--xml_message_type",
                 "invalid",
@@ -94,7 +94,7 @@ def test_main_with_invalid_xml_message_type(self):
 
     def test_main_with_invalid_xml_template_file(self):
         result = self.runner.invoke(
-            main,
+            cli,
             [
                 "--xml_message_type",
                 self.xml_message_type,
@@ -113,7 +113,7 @@ def test_main_with_invalid_xml_template_file(self):
 
     def test_main_with_invalid_xsd_template_file(self):
         result = self.runner.invoke(
-            main,
+            cli,
             [
                 "--xml_message_type",
                 self.xml_message_type,
@@ -132,7 +132,7 @@ def test_main_with_invalid_xsd_template_file(self):
 
     def test_main_with_invalid_data_file(self):
         result = self.runner.invoke(
-            main,
+            cli,
             [
                 "--xml_message_type",
                 self.xml_message_type,