Repository: google/turbinia
Pull request #487 (2019-12-29 20:19:41)

Proposed fix for issue #453
Creating as normal PR to see the test results. Let me know if a function-only implementation, such as the one in state_manager is a better fit.

jorlamd: Changed approach to that of the state manager to avoid code duplication in instantiating the individual clients across the code base, and because a factory class is less useful if we have the client definitions and the factory code in a single place. Changed the name of TurbiniaClient to BaseTurbiniaClient, as it is used as the base class for TurbiniaCeleryWorker and TurbiniaCeleryClient.
jorlamd: Thanks for reviewing. All proposed changes make sense to me. Let me know if you see anything further that we haven't addressed. Will do the same.



Patch:
diff --git a/turbinia/client.py b/turbinia/client.py
index cde7c3a..cc298dc 100644
--- a/turbinia/client.py
+++ b/turbinia/client.py
@@ -95,6 +95,24 @@ log = logging.getLogger('turbinia')
 logger.setup()
 
 
+def get_turbinia_client(run_local=False):
+  """Return Turbinia client based on config.
+
+  Returns:
+    Initialized BaseTurbiniaClient or TurbiniaCeleryClient object.
+  """
+  config.LoadConfig()
+  # pylint: disable=no-else-return
+  if config.TASK_MANAGER.lower() == 'psq':
+    return BaseTurbiniaClient(run_local=run_local)
+  elif config.TASK_MANAGER.lower() == 'celery':
+    return TurbiniaCeleryClient(run_local=run_local)
+  else:
+    msg = 'Task Manager type "{0:s}" not implemented'.format(
+        config.TASK_MANAGER)
+    raise TurbiniaException(msg)
+
+
 def check_dependencies(dependencies):
   """Checks system dependencies.
 
@@ -233,7 +251,7 @@ class TurbiniaStats(object):
         self.description, self.count, self.min, self.mean, self.max)
 
 
-class TurbiniaClient(object):
+class BaseTurbiniaClient(object):
   """Client class for Turbinia.
 
   Attributes:
@@ -748,7 +766,7 @@ class TurbiniaClient(object):
     return 'Closed Task IDs: %s' % response.get('result')
 
 
-class TurbiniaCeleryClient(TurbiniaClient):
+class TurbiniaCeleryClient(BaseTurbiniaClient):
   """Client class for Turbinia (Celery).
 
   Overriding some things specific to Celery operation.
@@ -757,8 +775,8 @@ class TurbiniaCeleryClient(TurbiniaClient):
     redis (RedisStateManager): Redis datastore object
   """
 
-  def __init__(self, *_, **__):
-    super(TurbiniaCeleryClient, self).__init__()
+  def __init__(self, *args, **kwargs):
+    super(TurbiniaCeleryClient, self).__init__(*args, **kwargs)
     self.redis = RedisStateManager()
 
   def send_request(self, request):
@@ -818,7 +836,7 @@ class TurbiniaServer(object):
     self.task_manager.add_evidence(evidence_)
 
 
-class TurbiniaCeleryWorker(TurbiniaClient):
+class TurbiniaCeleryWorker(BaseTurbiniaClient):
   """Turbinia Celery Worker class.
 
   Attributes:
diff --git a/turbinia/turbiniactl.py b/turbinia/turbiniactl.py
index 09e334b..8ba90d1 100644
--- a/turbinia/turbiniactl.py
+++ b/turbinia/turbiniactl.py
@@ -404,7 +404,7 @@ def main():
   # the config until after we parse the args so that we can use those arguments
   # to point to config paths.
   from turbinia import notify
-  from turbinia.client import TurbiniaClient
+  from turbinia import client as TurbiniaClientProvider
   from turbinia.client import TurbiniaCeleryClient
   from turbinia.client import TurbiniaServer
   from turbinia.client import TurbiniaCeleryWorker
@@ -454,15 +454,9 @@ def main():
           'Cannot open file {0:s} [{1!s}]'.format(args.filter_patterns_file, e))
 
   # Create Client object
+  client = None
   if args.command not in ('psqworker', 'server'):
-    if config.TASK_MANAGER.lower() == 'celery':
-      client = TurbiniaCeleryClient()
-    elif args.run_local:
-      client = TurbiniaClient(run_local=True)
-    else:
-      client = TurbiniaClient()
-  else:
-    client = None
+    client = TurbiniaClientProvider.get_turbinia_client(args.run_local)
 
   # Make sure run_local flags aren't conflicting with other server/client flags
   server_flags_set = args.server or args.command == 'server'
diff --git a/turbinia/workers/finalize_request.py b/turbinia/workers/finalize_request.py
index bb90463..2750edc 100644
--- a/turbinia/workers/finalize_request.py
+++ b/turbinia/workers/finalize_request.py
@@ -40,8 +40,8 @@ class FinalizeRequestTask(TurbiniaTask):
         TurbiniaTaskResult: Task execution results.
     """
     # Doing a delayed import to avoid circular dependencies.
-    from turbinia.client import TurbiniaClient
-    client = TurbiniaClient()
+    from turbinia import client as TurbiniaClientProvider
+    client = TurbiniaClientProvider.get_turbinia_client()
 
     report_file = os.path.join(
         self.tmp_dir, 'final_turbinia_report_{0:s}.md'.format(self.id))


Test patch:
diff --git a/turbinia/client_test.py b/turbinia/client_test.py
index b27626f..ebd2b2d 100644
--- a/turbinia/client_test.py
+++ b/turbinia/client_test.py
@@ -28,7 +28,7 @@ import textwrap
 import mock
 
 from turbinia import config
-from turbinia.client import TurbiniaClient
+from turbinia import client as TurbiniaClientProvider
 from turbinia.client import TurbiniaServer
 from turbinia.client import TurbiniaStats
 from turbinia.client import TurbiniaPsqWorker
@@ -205,7 +205,7 @@ class TestTurbiniaClient(unittest.TestCase):
   def testTurbiniaClientInit(self, _, __):
     """Basic test for client."""
     config.LoadConfig()
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     self.assertTrue(hasattr(client, 'task_manager'))
 
   @mock.patch('turbinia.client.GoogleCloudFunction.ExecuteFunction')
@@ -221,7 +221,7 @@ class TestTurbiniaClient(unittest.TestCase):
     gcf_result = json.dumps(gcf_result)
     function_return = {'result': gcf_result}
     mock_cloud_function.return_value = function_return
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     task_data = client.get_task_data('inst', 'proj', 'reg')
     # get_task_data() converts this back into a timedelta(). We returned it
     # seconds from the GCF function call because that is what it is stored in
@@ -235,7 +235,7 @@ class TestTurbiniaClient(unittest.TestCase):
   def testTurbiniaClientGetTaskDataNoResults(self, _, __, mock_cloud_function):
     """Test for exception after empty results from cloud functions."""
     mock_cloud_function.return_value = {}
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     self.assertRaises(
         TurbiniaException, client.get_task_data, "inst", "proj", "reg")
 
@@ -246,7 +246,7 @@ class TestTurbiniaClient(unittest.TestCase):
       self, _, __, mock_cloud_function):
     """Test for exception after bad json results from cloud functions."""
     mock_cloud_function.return_value = {'result': None}
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     self.assertRaises(
         TurbiniaException, client.get_task_data, "inst", "proj", "reg")
 
@@ -255,7 +255,7 @@ class TestTurbiniaClient(unittest.TestCase):
   @mock.patch('turbinia.state_manager.get_state_manager')
   def testClientFormatTaskStatistics(self, _, __, ___):
     """Tests format_task_statistics() report output."""
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     client.get_task_data = mock.MagicMock()
     client.get_task_data.return_value = self.task_data
     stats_report = client.format_task_statistics('inst', 'proj', 'reg')
@@ -267,7 +267,7 @@ class TestTurbiniaClient(unittest.TestCase):
   @mock.patch('turbinia.state_manager.get_state_manager')
   def testClientFormatTaskStatisticsCsv(self, _, __, ___):
     """Tests format_task_statistics() CSV report output."""
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     client.get_task_data = mock.MagicMock()
     client.get_task_data.return_value = self.task_data
     stats_report = client.format_task_statistics(
@@ -280,7 +280,7 @@ class TestTurbiniaClient(unittest.TestCase):
   @mock.patch('turbinia.state_manager.get_state_manager')
   def testClientGetTaskStatistics(self, _, __, ___):
     """Tests get_task_statistics() basic functionality."""
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     client.get_task_data = mock.MagicMock()
     client.get_task_data.return_value = self.task_data
     task_stats = client.get_task_statistics('inst', 'proj', 'reg')
@@ -313,7 +313,7 @@ class TestTurbiniaClient(unittest.TestCase):
   @mock.patch('turbinia.state_manager.get_state_manager')
   def testClientFormatTaskStatus(self, _, __, ___):
     """Tests format_task_status() with empty report_priority."""
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     client.get_task_data = mock.MagicMock()
     self.task_data[0]['report_priority'] = None
     self.task_data[1]['report_priority'] = ''
@@ -327,7 +327,7 @@ class TestTurbiniaClient(unittest.TestCase):
   @mock.patch('turbinia.state_manager.get_state_manager')
   def testClientFormatTaskStatusShortReport(self, _, __, ___):
     """Tests format_task_status() has valid output with short report."""
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     client.get_task_data = mock.MagicMock()
     client.get_task_data.return_value = self.task_data
     result = client.format_task_status('inst', 'proj', 'reg')
@@ -338,7 +338,7 @@ class TestTurbiniaClient(unittest.TestCase):
   @mock.patch('turbinia.state_manager.get_state_manager')
   def testClientFormatTaskStatusFullReport(self, _, __, ___):
     """Tests format_task_status() has valid output with full report."""
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     client.get_task_data = mock.MagicMock()
     client.get_task_data.return_value = self.task_data
     result = client.format_task_status('inst', 'proj', 'reg', full_report=True)
@@ -349,7 +349,7 @@ class TestTurbiniaClient(unittest.TestCase):
   @mock.patch('turbinia.state_manager.get_state_manager')
   def testClientFormatTaskStatusFiles(self, _, __, ___):
     """Tests format_task_status() has valid output with report and files."""
-    client = TurbiniaClient()
+    client = TurbiniaClientProvider.get_turbinia_client()
     client.get_task_data = mock.MagicMock()
     client.get_task_data.return_value = self.task_data
     result = client.format_task_status(