From 96c52a032821765474f7894b84ac2614428ee4f8 Mon Sep 17 00:00:00 2001
From: Jake Howard <git@theorangeone.net>
Date: Wed, 12 Dec 2018 18:47:40 +0000
Subject: [PATCH] Test worker stop/start works

---
 catfish/__main__.py        | 21 ++++++++++++++-------
 catfish/worker/__init__.py | 16 ++++++++++++++++
 tests/__init__.py          |  7 +++++++
 tests/test_main.py         | 38 ++++++++++++++++++++++++++++++++++++--
 4 files changed, 73 insertions(+), 9 deletions(-)

diff --git a/catfish/__main__.py b/catfish/__main__.py
index acfc440..cd6725e 100644
--- a/catfish/__main__.py
+++ b/catfish/__main__.py
@@ -1,8 +1,10 @@
+import os
+import sys
+
 import click
 import daemonize
 
 from catfish import __version__, worker
-from catfish.utils.processes import terminate_processes
 
 
 @click.group()
@@ -13,26 +15,31 @@ def cli():
 
 @cli.command()
 @click.option("--no-fork", is_flag=True)
-def start(no_fork):
+@click.pass_context
+def start(ctx, no_fork):
+    if worker.is_running():
+        ctx.fail("Worker already running")
     daemon = daemonize.Daemonize(
         "catfish", worker.PID_FILE, worker.main, foreground=no_fork
     )
     try:
+        # HACK: Temporary hack until https://github.com/thesharp/daemonize/pull/70 is solved
+        os._exit = sys.exit
         daemon.start()
     except SystemExit:
-        return
+        worker.wait_for_worker()
     proc = worker.get_running_process()
     click.echo("Worker started with pid {}".format(proc.pid))
 
 
 @cli.command()
-def stop():
+@click.pass_context
+def stop(ctx):
     if not worker.is_running():
-        click.echo("Worker not running")
-        return
+        ctx.fail("Worker not running")
     proc = worker.get_running_process()
     click.echo("Terminating process {}".format(proc.pid))
-    terminate_processes([proc])
+    worker.stop_worker()
 
 
 if __name__ == "__main__":
diff --git a/catfish/worker/__init__.py b/catfish/worker/__init__.py
index a79df02..11eaa1b 100644
--- a/catfish/worker/__init__.py
+++ b/catfish/worker/__init__.py
@@ -2,6 +2,7 @@ import time
 import os
 import tempfile
 import psutil
+from catfish.utils.processes import terminate_processes
 
 
 PID_FILE = os.path.join(tempfile.gettempdir(), "catfish.pid")
@@ -17,6 +18,21 @@ def get_running_process() -> psutil.Process:
         return psutil.Process(int(f.read()))
 
 
+def wait_for_worker():
+    while True:
+        try:
+            get_running_process()
+            return
+        except ValueError:
+            pass
+        time.sleep(0.1)
+
+
+def stop_worker():
+    if is_running():
+        terminate_processes([get_running_process()])
+
+
 def main():
     while True:
         time.sleep(1)
diff --git a/tests/__init__.py b/tests/__init__.py
index 6c83a8d..4cea90f 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -1,9 +1,16 @@
 from unittest import TestCase
 from click.testing import CliRunner
 from catfish.__main__ import cli
+from catfish import worker
+import functools
 
 
 class BaseTestCase(TestCase):
     def setUp(self):
+        worker.stop_worker()
         self.cli_runner = CliRunner()
         self.cli = cli
+        self.run_cli = functools.partial(self.cli_runner.invoke, self.cli)
+
+    def tearDown(self):
+        worker.stop_worker()
diff --git a/tests/test_main.py b/tests/test_main.py
index 9173b0c..a46a24f 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -1,10 +1,44 @@
-from catfish import __version__
+from catfish import __version__, worker
 from tests import BaseTestCase
 
 
 class MainCLITestCase(BaseTestCase):
     def test_version(self):
-        result = self.cli_runner.invoke(self.cli, ["--version"])
+        result = self.run_cli(["--version"])
         self.assertEqual(result.exit_code, 0)
         self.assertIn(__version__, result.output)
         self.assertIn("catfish", result.output)
+
+
+class WorkerControlTestCase(BaseTestCase):
+    def test_starts_worker(self):
+        result = self.run_cli(["start"])
+        self.assertEqual(result.exit_code, 0)
+        self.assertTrue(worker.is_running())
+        worker_process = worker.get_running_process()
+        self.assertIn(str(worker_process.pid), result.output)
+
+    def test_stops_worker(self):
+        result = self.run_cli(["start"])
+        self.assertEqual(result.exit_code, 0)
+        self.assertTrue(worker.is_running())
+        worker_pid = worker.get_running_process().pid
+        result = self.run_cli(["stop"])
+        self.assertEqual(result.exit_code, 0)
+        self.assertIn(str(worker_pid), result.output)
+        self.assertFalse(worker.is_running())
+
+    def test_starting_when_already_running(self):
+        result = self.run_cli(["start"])
+        self.assertEqual(result.exit_code, 0)
+        self.assertTrue(worker.is_running())
+        result = self.run_cli(["start"])
+        self.assertEqual(result.exit_code, 2)
+        self.assertIn("Worker already running", result.output)
+        self.assertTrue(worker.is_running())
+
+    def test_stop_when_not_running(self):
+        self.assertFalse(worker.is_running())
+        result = self.run_cli(["stop"])
+        self.assertEqual(result.exit_code, 2)
+        self.assertIn("Worker not running", result.output)