From df7c0b8e32c448b1bd3a2a9471e17b181aabdac3 Mon Sep 17 00:00:00 2001
From: James Kafader <jkafader@archive.org>
Date: Tue, 3 Oct 2017 14:38:31 -0700
Subject: [PATCH] added tests for purging stale services and minimal tests for
 command line tool

---
 doublethink/cli.py   |  3 ++-
 tests/test_cli.py    | 27 ++++++++++++++++++++++++---
 tests/test_svcreg.py | 10 ++++++++++
 3 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/doublethink/cli.py b/doublethink/cli.py
index 517e52e..fab424c 100644
--- a/doublethink/cli.py
+++ b/doublethink/cli.py
@@ -54,4 +54,5 @@ def purge_stale_services(argv=None):
 
     rethinker = doublethink.Rethinker(servers=args.servers, db=args.database)
     registry = doublethink.services.ServiceRegistry(rethinker)
-    registry.purge_stale_services()
\ No newline at end of file
+    registry.purge_stale_services()
+    sys.exit(0)
\ No newline at end of file
diff --git a/tests/test_cli.py b/tests/test_cli.py
index fb7f43a..45c88b4 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -17,14 +17,25 @@ limitations under the License.
 '''
 
 import doublethink
+import doublethink.cli
 import logging
 import sys
 import pytest
 import rethinkdb as r
+import pkg_resources
 
 logging.basicConfig(stream=sys.stderr, level=logging.INFO,
         format="%(asctime)s %(process)d %(levelname)s %(threadName)s %(name)s.%(funcName)s(%(filename)s:%(lineno)d) %(message)s")
 
+class RethinkerForTesting(doublethink.Rethinker):
+    def __init__(self, *args, **kwargs):
+        super(RethinkerForTesting, self).__init__(*args, **kwargs)
+
+    def _random_server_connection(self):
+        self.last_conn = super(RethinkerForTesting, self)._random_server_connection()
+        # logging.info("self.last_conn=%s", self.last_conn)
+        return self.last_conn
+
 @pytest.fixture(scope="module")
 def rr():
     rr = RethinkerForTesting()
@@ -37,6 +48,16 @@ def rr():
     assert result["dbs_created"] == 1
     return RethinkerForTesting(db="doublethink_test_db")
 
-def test_cli(rr):
-    print(rr)
-    doublethink.cli.purge_stale_services(['test'])
+def test_cli(capsys, rr):
+    entrypoint = pkg_resources.get_entry_map(
+            'doublethink')['console_scripts']['purge-stale-services']
+    callable = entrypoint.resolve()
+    with pytest.raises(SystemExit) as exit:
+        callable(['purge-stale-services'])
+    print(dir(exit))
+    assert exit.value.code != 0
+    out, err = capsys.readouterr()
+    with pytest.raises(SystemExit) as exit:
+        callable(['purge-stale-services', '-d', 'test'])
+    assert exit.value.code == 0
+    out, err = capsys.readouterr()
diff --git a/tests/test_svcreg.py b/tests/test_svcreg.py
index aae5e22..380d115 100644
--- a/tests/test_svcreg.py
+++ b/tests/test_svcreg.py
@@ -258,3 +258,13 @@ def test_svcreg_heartbeat_server_down(rr):
     assert not 'host' in svc0
     assert not 'pid' in svc0
 
+def test_purge_stale_services(rr):
+    rr.table('services').delete().run()
+    rr.table('services').insert({ 'id': 'old-service', "last_heartbeat": r.now(), 'ttl': 0.4 }).run()
+    time.sleep(1)
+    rr.table('services').insert({ 'id': 'new-service', "last_heartbeat": r.now(), 'ttl': 0.4 }).run()
+    svcreg = doublethink.ServiceRegistry(rr)
+    assert rr.table('services').count().run() == 2
+    svcreg.purge_stale_services()
+    assert rr.table('services').count().run() == 1
+    rr.table('services').delete().run()