From 243c21db00e436326567b866c0f78b3c346d8946 Mon Sep 17 00:00:00 2001
From: ouhammmourachid <rachidouhammou21@gmail.com>
Date: Mon, 6 Nov 2023 15:27:08 +0100
Subject: [PATCH] Update supported file format for script saving to .mmd or
 .mermaid #40

---
 mermaid/graph/__init__.py   |  8 +++-----
 mermaid/tests/test_graph.py | 27 ++++++++++++++++-----------
 mermaid/tests/test_utils.py |  4 ++--
 test-graph-file.txt         |  5 +++++
 4 files changed, 26 insertions(+), 18 deletions(-)
 create mode 100644 test-graph-file.txt

diff --git a/mermaid/graph/__init__.py b/mermaid/graph/__init__.py
index a520ebd..261dee3 100644
--- a/mermaid/graph/__init__.py
+++ b/mermaid/graph/__init__.py
@@ -15,16 +15,14 @@ class Graph:
     examples:
     >>> from mermaid.graph import Graph
     >>> grph = Graph("title","...")
-    >>> # save the script in txt file for other uses
-    >>> graph.save() # save in file with name of title and extention .txt.
-    >>> graph.svae('./file-name.txt') # save in specific file.
     """
     title: str
     script: str
 
     def save(self, path: Optional[Path] = None) -> None:
-        file_path: Path = path if path else Path(f'./{self.title}.txt')
-
+        file_path: Path = path if path else Path(f'./{self.title}.mmd')
+        if file_path.suffix not in ['.mmd', '.mermaid']:
+            raise ValueError("File extension must be '.mmd' or '.mermaid'")
         with open(file_path, 'w') as file:
             file.write(self.script)
 
diff --git a/mermaid/tests/test_graph.py b/mermaid/tests/test_graph.py
index 64e6b7f..c5bbe41 100644
--- a/mermaid/tests/test_graph.py
+++ b/mermaid/tests/test_graph.py
@@ -17,22 +17,27 @@ def setUp(self) -> None:
 
     def test_save_graph_without_path(self) -> None:
         self.graph_test.save()
-        self.assertTrue(Path.exists(Path('./test-graph.txt')))
+        self.assertTrue(Path.exists(Path('./test-graph.mmd')))
 
     def test_save_graph_with_path(self) -> None:
-        self.graph_test.save(Path('./test-graph-file.txt'))
-        self.assertTrue(Path.exists(Path('./test-graph-file.txt')))
+        self.graph_test.save(Path('./test-graph-file.mmd'))
+        self.assertTrue(Path.exists(Path('./test-graph-file.mmd')))
 
     def test_save_graph_name_contain_space(self) -> None:
         self.graph_test.title = 'test graph file'
-        self.graph_test.save(Path('./test graph file.txt'))
-        self.assertTrue(Path.exists(Path('./test graph file.txt')))
+        self.graph_test.save(Path('./test graph file.mmd'))
+        self.assertTrue(Path.exists(Path('./test graph file.mmd')))
+
+    def test_save_should_raise_valueerror(self):
+        self.graph_test.title = 'file-name'
+        with self.assertRaises(ValueError):
+            self.graph_test.save(Path('./file-name.txt'))
 
     def tearDown(self) -> None:
-        if os.path.exists('./test-graph.txt'):
-            os.remove('./test-graph.txt')
-        if os.path.exists('./test-graph-file.txt'):
-            os.remove('./test-graph-file.txt')
-        if os.path.exists('./test graph file.txt'):
-            os.remove('./test graph file.txt')
+        if os.path.exists('./test-graph.mmd'):
+            os.remove('./test-graph.mmd')
+        if os.path.exists('./test-graph-file.mmd'):
+            os.remove('./test-graph-file.mmd')
+        if os.path.exists('./test graph file.mmd'):
+            os.remove('./test graph file.mmd')
         return super().tearDown()
diff --git a/mermaid/tests/test_utils.py b/mermaid/tests/test_utils.py
index 162c784..126b137 100644
--- a/mermaid/tests/test_utils.py
+++ b/mermaid/tests/test_utils.py
@@ -7,7 +7,7 @@
 
 class TestUtils(unittest.TestCase):
     def setUp(self) -> None:
-        self.file_test: str = './test-graph.txt'
+        self.file_test: str = './test-graph.mmd'
         with open(self.file_test, 'w') as file:
             file.write("""graph TD;
     A-->B;
@@ -18,7 +18,7 @@ def setUp(self) -> None:
 
     def test_read_file_should_raise_error(self):
         with self.assertRaises(FileNotFoundError):
-            md.load('./file.txt')
+            md.load('./file.mmd')
 
     def test_read_file(self):
         graph: Graph = md.load(self.file_test)
diff --git a/test-graph-file.txt b/test-graph-file.txt
new file mode 100644
index 0000000..d5bf6cb
--- /dev/null
+++ b/test-graph-file.txt
@@ -0,0 +1,5 @@
+graph TD;
+    A-->B;
+    A-->C;
+    B-->D;
+    C-->D;