Refactor and clean up and add tests

Author: ptrlrd

main.py

@@ -242,12 +242,12 @@ def schedule_post(text: str, post_time: str):
     typer.echo(f"Post scheduled for {post_time} with text: '{text}'")
 
 @app.command()
-def create_draft(text: str):
-    """
-    Create a draft with text.
-    """
-    if os.path.exists(DRAFTS_FILE):
-        with open(DRAFTS_FILE, 'r') as file:
+def create_draft(text: str, drafts_file: str = DRAFTS_FILE):
+    '''
+    Create a draft with the given text and save it to the drafts file.
+    '''
+    if os.path.exists(drafts_file):
+        with open(drafts_file, 'r') as file:
             drafts = json.load(file)
     else:
         drafts = []
@@ -261,21 +261,21 @@ def create_draft(text: str):
     }
     drafts.append(draft)
 
-    with open(DRAFTS_FILE, 'w') as file:
+    with open(drafts_file, 'w') as file:
         json.dump(drafts, file, indent=4)
 
     typer.echo(f"Draft created with ID: {next_id}")
 
 @app.command()
-def get_drafts():
-    """
-    Retrieve all drafts.
-    """
-    if not os.path.exists(DRAFTS_FILE):
+def get_drafts(drafts_file: str = DRAFTS_FILE):
+    '''
+    Get all drafts from the drafts file.
+    '''
+    if not os.path.exists(drafts_file):
         typer.echo("No drafts found.")
         return
 
-    with open(DRAFTS_FILE, 'r') as file:
+    with open(drafts_file, 'r') as file:
         drafts = json.load(file)
 
     table = Table(title="Drafts")
@@ -293,28 +293,28 @@ def get_drafts():
     console.print(table)
 
 @app.command()
-def send_draft(draft_id: int):
-    """
-    Sends a draft with a specific ID and removes it from drafts.
-    """
-    if not os.path.exists(DRAFTS_FILE):
+def send_draft(draft_id: int, drafts_file: str = DRAFTS_FILE):
+    '''
+    Send a draft with the given ID and remove it from the drafts file.
+    '''
+    if not os.path.exists(drafts_file):
         typer.echo("No drafts found.")
-        return
+        raise typer.Exit(1)
 
-    with open(DRAFTS_FILE, 'r') as file:
+    with open(drafts_file, 'r') as file:
         drafts = json.load(file)
 
     draft = next((draft for draft in drafts if draft['id'] == draft_id), None)
 
     if draft is None:
         typer.echo(f"Draft with ID {draft_id} not found.")
-        return
+        raise typer.Exit(1)
 
     create_text_post(draft['text'])
 
     drafts = [draft for draft in drafts if draft['id'] != draft_id]
 
-    with open(DRAFTS_FILE, 'w') as file:
+    with open(drafts_file, 'w') as file:
         json.dump(drafts, file, indent=4)
 
     typer.echo(f"Draft with ID {draft_id} sent and removed from drafts.")

test_main.py

@@ -0,0 +1,73 @@
+import unittest
+import os
+import json
+from unittest.mock import patch
+from typer.testing import CliRunner
+from main import app
+
+TEST_DRAFTS_FILE = "test-drafts.json"
+
+class TestThreadsCLI(unittest.TestCase):
+    def setUp(self):
+        self.runner = CliRunner()
+
+    def tearDown(self):
+        if os.path.exists(TEST_DRAFTS_FILE):
+            os.remove(TEST_DRAFTS_FILE)
+
+    def test_get_profile(self):
+        result = self.runner.invoke(app, ["get-profile"])
+        self.assertEqual(result.exit_code, 0)
+        self.assertIn("Profile", result.stdout)
+
+    def test_get_recent_posts(self):
+        result = self.runner.invoke(app, ["get-recent-posts", "--limit", "5"])
+        self.assertEqual(result.exit_code, 0)
+        self.assertIn("Recent Posts", result.stdout)
+
+    def test_create_text_post(self):
+        with patch("main.create_post") as mock_create_post:
+            mock_create_post.return_value = {"id": "123"}
+            result = self.runner.invoke(app, ["create-text-post", "Test post"])
+            self.assertEqual(result.exit_code, 0)
+            self.assertIn("Post created with ID: 123", result.stdout)
+
+    def test_create_draft(self):
+        result = self.runner.invoke(app, ["create-draft", "Test draft", "--drafts-file", TEST_DRAFTS_FILE])
+        self.assertEqual(result.exit_code, 0)
+        self.assertTrue(os.path.exists(TEST_DRAFTS_FILE))
+        with open(TEST_DRAFTS_FILE, "r") as file:
+            drafts = json.load(file)
+            self.assertEqual(len(drafts), 1)
+            self.assertEqual(drafts[0]["text"], "Test draft")
+
+    def test_get_drafts(self):
+        self.runner.invoke(app, ["create-draft", "Test draft 1", "--drafts-file", TEST_DRAFTS_FILE])
+        self.runner.invoke(app, ["create-draft", "Test draft 2", "--drafts-file", TEST_DRAFTS_FILE])
+        result = self.runner.invoke(app, ["get-drafts", "--drafts-file", TEST_DRAFTS_FILE])
+        self.assertEqual(result.exit_code, 0)
+        self.assertIn("Test draft 1", result.stdout)
+        self.assertIn("Test draft 2", result.stdout)
+
+    def test_send_draft(self):
+        self.runner.invoke(app, ["create-draft", "Test draft", "--drafts-file", TEST_DRAFTS_FILE])
+        with open(TEST_DRAFTS_FILE, "r") as file:
+            drafts = json.load(file)
+            draft_id = drafts[0]["id"]
+
+        # Test sending an existing draft
+        with patch("main.create_text_post") as mock_create_text_post:
+            result = self.runner.invoke(app, ["send-draft", str(draft_id), "--drafts-file", TEST_DRAFTS_FILE])
+            self.assertEqual(result.exit_code, 0)
+            self.assertIn(f"Draft with ID {draft_id} sent and removed from drafts.", result.stdout)
+        with open(TEST_DRAFTS_FILE, "r") as file:
+            drafts = json.load(file)
+            self.assertEqual(len(drafts), 0)
+
+        # Test sending a non-existing draft
+        result = self.runner.invoke(app, ["send-draft", "999", "--drafts-file", TEST_DRAFTS_FILE])
+        self.assertEqual(result.exit_code, 1)
+        self.assertIn("Draft with ID 999 not found.", result.stdout)
+
+if __name__ == "__main__":
+    unittest.main()