Fix tests, and issues caught by tests

Fix some issues caught by tests due to changes introducing the hand-off of the
partial correlations computations to an external process

Fix some issues due to the changes that introduce context managers for
database connections

Update some tests to take the above two changes into consideration

9 files changed, 57 insertions, 50 deletions

diff --git a/gn3/api/correlation.py b/gn3/api/correlation.py
index f2ac4d7..67897e8 100644
--- a/gn3/api/correlation.py
+++ b/gn3/api/correlation.py
@@ -112,7 +112,7 @@ def partial_correlation():
 
         return reduce(__field_errors__(request_data), fields, errors)
 
-    args = json.loads(request.get_json())
+    args = request.get_json()
     request_errors = __errors__(
         args, ("primary_trait", "control_traits", "target_db", "method"))
     if request_errors:
diff --git a/gn3/commands.py b/gn3/commands.py
index 29e3df2..0e0b53c 100644
--- a/gn3/commands.py
+++ b/gn3/commands.py
@@ -84,11 +84,12 @@ Returns the name of the specific redis hash for the specific task.
                  f"{str(uuid4())}")
     conn.rpush(job_queue, unique_id)
     for key, value in {
-            "cmd": json.dumps(cmd), "result": "", "status": "queued",
-            "env": json.dumps(env)}.items():
+            "cmd": json.dumps(cmd), "result": "", "status": "queued"}.items():
         conn.hset(name=unique_id, key=key, value=value)
     if email:
         conn.hset(name=unique_id, key="email", value=email)
+    if env:
+        conn.hset(name=unique_id, key="env", value=json.dumps(env))
     return unique_id
 
 
diff --git a/gn3/computations/diff.py b/gn3/computations/diff.py
index af02f7f..0b6edd6 100644
--- a/gn3/computations/diff.py
+++ b/gn3/computations/diff.py
@@ -6,7 +6,7 @@ from gn3.commands import run_cmd
 
 def generate_diff(data: str, edited_data: str) -> Optional[str]:
     """Generate the diff between 2 files"""
-    results = run_cmd(f"diff {data} {edited_data}", success_codes=(1, 2))
+    results = run_cmd(f'"diff {data} {edited_data}"', success_codes=(1, 2))
     if results.get("code", -1) > 0:
         return results.get("output")
     return None
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
index e1d1c37..be927a4 100644
--- a/tests/integration/conftest.py
+++ b/tests/integration/conftest.py
@@ -19,4 +19,6 @@ def client():
 def db_conn():
     """Create a db connection fixture for tests"""
     ## Update this to use temp db once that is in place
-    return database_connector()[0]
+    conn = database_connector()
+    yield conn
+    conn.close()
diff --git a/tests/integration/test_correlation.py b/tests/integration/test_correlation.py
index cf63c17..d52ab01 100644
--- a/tests/integration/test_correlation.py
+++ b/tests/integration/test_correlation.py
@@ -71,7 +71,9 @@ class CorrelationIntegrationTest(TestCase):
 
         mock_compute_corr.return_value = []
 
-        database_connector.return_value = (mock.Mock(), mock.Mock())
+        database_connector.return_value = mock.Mock()
+        database_connector.return_value.__enter__ = mock.Mock()
+        database_connector.return_value.__exit__ = mock.Mock()
 
         post_data = {"1426678_at": "68031",
                      "1426679_at": "68036",
diff --git a/tests/integration/test_partial_correlations.py b/tests/integration/test_partial_correlations.py
index 7f9ff30..d249b42 100644
--- a/tests/integration/test_partial_correlations.py
+++ b/tests/integration/test_partial_correlations.py
@@ -1,4 +1,6 @@
 """Test partial correlations"""
+from unittest import mock
+
 import pytest
 
 from gn3.computations.partial_correlations import partial_correlations_entry
@@ -97,43 +99,46 @@ def test_partial_correlation_api_with_missing_request_data(client, post_data):
 @pytest.mark.parametrize(
     "post_data",
     ({# ProbeSet
-        "primary_trait": {"dataset": "a_dataset", "name": "a_name"},
+        "primary_trait": {"dataset": "a_dataset", "trait_name": "a_name"},
         "control_traits": [
-            {"dataset": "a_dataset", "name": "a_name"},
-            {"dataset": "a_dataset2", "name": "a_name2"}],
+            {"dataset": "a_dataset", "trait_name": "a_name"},
+            {"dataset": "a_dataset2", "trait_name": "a_name2"}],
         "method": "a_method",
         "target_db": "a_db"
     }, {# Publish
-        "primary_trait": {"dataset": "a_Publish_dataset", "name": "a_name"},
+        "primary_trait": {
+            "dataset": "a_Publish_dataset", "trait_name": "a_name"},
         "control_traits": [
-            {"dataset": "a_dataset", "name": "a_name"},
-            {"dataset": "a_dataset2", "name": "a_name2"}],
+            {"dataset": "a_dataset", "trait_name": "a_name"},
+            {"dataset": "a_dataset2", "trait_name": "a_name2"}],
         "method": "a_method",
         "target_db": "a_db"
     }, {# Geno
-        "primary_trait": {"dataset": "a_Geno_dataset", "name": "a_name"},
+        "primary_trait": {"dataset": "a_Geno_dataset", "trait_name": "a_name"},
         "control_traits": [
-            {"dataset": "a_dataset", "name": "a_name"},
-            {"dataset": "a_dataset2", "name": "a_name2"}],
+            {"dataset": "a_dataset", "trait_name": "a_name"},
+            {"dataset": "a_dataset2", "trait_name": "a_name2"}],
         "method": "a_method",
         "target_db": "a_db"
     }, {# Temp
-        "primary_trait": {"dataset": "a_Temp_dataset", "name": "a_name"},
+        "primary_trait": {"dataset": "a_Temp_dataset", "trait_name": "a_name"},
         "control_traits": [
-            {"dataset": "a_dataset", "name": "a_name"},
-            {"dataset": "a_dataset2", "name": "a_name2"}],
+            {"dataset": "a_dataset", "trait_name": "a_name"},
+            {"dataset": "a_dataset2", "trait_name": "a_name2"}],
         "method": "a_method",
         "target_db": "a_db"
     }))
-def test_partial_correlation_api_with_non_existent_primary_traits(client, post_data):
+def test_partial_correlation_api_with_non_existent_primary_traits(
+        client, post_data, mocker):
     """
     Check that the system responds appropriately in the case where the user
     makes a request with a non-existent primary trait.
     """
+    mocker.patch("gn3.api.correlation.redis.Redis", mock.MagicMock())
     response = client.post("/api/correlation/partial", json=post_data)
     assert (
-        response.status_code == 404 and response.is_json and
-        response.json.get("status") != "error")
+        response.status_code == 200 and response.is_json and
+        response.json.get("status") == "success")
 
 @pytest.mark.integration_test
 @pytest.mark.slow
@@ -141,41 +146,43 @@ def test_partial_correlation_api_with_non_existent_primary_traits(client, post_d
     "post_data",
     ({# ProbeSet
         "primary_trait": {
-            "dataset": "UCLA_BXDBXH_CARTILAGE_V2", "name": "ILM103710672"},
+            "dataset": "UCLA_BXDBXH_CARTILAGE_V2",
+            "trait_name": "ILM103710672"},
         "control_traits": [
-            {"dataset": "a_dataset", "name": "a_name"},
-            {"dataset": "a_dataset2", "name": "a_name2"}],
+            {"dataset": "a_dataset", "trait_name": "a_name"},
+            {"dataset": "a_dataset2", "trait_name": "a_name2"}],
         "method": "a_method",
         "target_db": "a_db"
     }, {# Publish
-        "primary_trait": {"dataset": "BXDPublish", "name": "BXD_12557"},
+        "primary_trait": {"dataset": "BXDPublish", "trait_name": "BXD_12557"},
         "control_traits": [
-            {"dataset": "a_dataset", "name": "a_name"},
-            {"dataset": "a_dataset2", "name": "a_name2"}],
+            {"dataset": "a_dataset", "trait_name": "a_name"},
+            {"dataset": "a_dataset2", "trait_name": "a_name2"}],
         "method": "a_method",
         "target_db": "a_db"
     }, {# Geno
-        "primary_trait": {"dataset": "AKXDGeno", "name": "D4Mit16"},
+        "primary_trait": {"dataset": "AKXDGeno", "trait_name": "D4Mit16"},
         "control_traits": [
-            {"dataset": "a_dataset", "name": "a_name"},
-            {"dataset": "a_dataset2", "name": "a_name2"}],
+            {"dataset": "a_dataset", "trait_name": "a_name"},
+            {"dataset": "a_dataset2", "trait_name": "a_name2"}],
         "method": "a_method",
         "target_db": "a_db"
     }
      # Temp -- the data in the database for these is ephemeral, making it
      #         difficult to test for this
      ))
-def test_partial_correlation_api_with_non_existent_control_traits(client, post_data):
+def test_partial_correlation_api_with_non_existent_control_traits(client, post_data, mocker):
     """
     Check that the system responds appropriately in the case where the user
     makes a request with a non-existent control traits.
 
     The code repetition here is on purpose - valuing clarity over succinctness.
     """
+    mocker.patch("gn3.api.correlation.redis.Redis", mock.MagicMock())
     response = client.post("/api/correlation/partial", json=post_data)
     assert (
-        response.status_code == 404 and response.is_json and
-        response.json.get("status") != "error")
+        response.status_code == 200 and response.is_json and
+        response.json.get("status") == "success")
 
 @pytest.mark.integration_test
 @pytest.mark.slow
diff --git a/tests/performance/perf_query.py b/tests/performance/perf_query.py
index c22dcf5..e534e9b 100644
--- a/tests/performance/perf_query.py
+++ b/tests/performance/perf_query.py
@@ -28,15 +28,13 @@ def timer(func):
 def query_executor(query: str,
                    fetch_all: bool = True):
     """function to execute a query"""
-    conn, _ = database_connector()
+    with database_connector() as conn:
+        with conn.cursor() as cursor:
+            cursor.execute(query)
 
-    with conn:
-        cursor = conn.cursor()
-        cursor.execute(query)
-
-        if fetch_all:
-            return cursor.fetchall()
-        return cursor.fetchone()
+            if fetch_all:
+                return cursor.fetchall()
+            return cursor.fetchone()
 
 
 def fetch_probeset_query(dataset_name: str):
diff --git a/tests/unit/test_commands.py b/tests/unit/test_commands.py
index e0efaf7..4dd8735 100644
--- a/tests/unit/test_commands.py
+++ b/tests/unit/test_commands.py
@@ -96,8 +96,7 @@ class TestCommands(unittest.TestCase):
     @pytest.mark.unit_test
     @mock.patch("gn3.commands.datetime")
     @mock.patch("gn3.commands.uuid4")
-    def test_queue_cmd_correct_calls_to_redis(self, mock_uuid4,
-                                              mock_datetime):
+    def test_queue_cmd_correct_calls_to_redis(self, mock_uuid4, mock_datetime):
         """Test that the cmd is queued properly"""
         mock_uuid4.return_value = 1234
         mock_datetime.now.return_value = datetime.fromisoformat('2021-02-12 '
@@ -112,7 +111,7 @@ class TestCommands(unittest.TestCase):
                                    job_queue="GN2::job-queue"),
                          actual_unique_id)
         mock_redis_conn.hset.assert_has_calls(
-            [mock.call(name=actual_unique_id, key="cmd", value="ls"),
+            [mock.call(name=actual_unique_id, key="cmd", value='"ls"'),
              mock.call(name=actual_unique_id, key="result", value=""),
              mock.call(name=actual_unique_id, key="status", value="queued")])
         mock_redis_conn.rpush.assert_has_calls(
@@ -139,23 +138,23 @@ class TestCommands(unittest.TestCase):
                                    email="me@me.com"),
                          actual_unique_id)
         mock_redis_conn.hset.assert_has_calls(
-            [mock.call(name=actual_unique_id, key="cmd", value="ls"),
+            [mock.call(name=actual_unique_id, key="cmd", value='"ls"'),
              mock.call(name=actual_unique_id, key="result", value=""),
              mock.call(name=actual_unique_id, key="status", value="queued"),
              mock.call(name=actual_unique_id, key="email", value="me@me.com")
-             ])
+             ], any_order=True)
         mock_redis_conn.rpush.assert_has_calls(
             [mock.call("GN2::job-queue", actual_unique_id)])
 
     @pytest.mark.unit_test
     def test_run_cmd_correct_input(self):
         """Test that a correct cmd is processed correctly"""
-        self.assertEqual(run_cmd("echo test"),
+        self.assertEqual(run_cmd('"echo test"'),
                          {"code": 0, "output": "test\n"})
 
     @pytest.mark.unit_test
     def test_run_cmd_incorrect_input(self):
         """Test that an incorrect cmd is processed correctly"""
-        result = run_cmd("echoo test")
+        result = run_cmd('"echoo test"')
         self.assertEqual(127, result.get("code"))
         self.assertIn("not found", result.get("output"))
diff --git a/tests/unit/test_db_utils.py b/tests/unit/test_db_utils.py
index 96ee68f..44d6703 100644
--- a/tests/unit/test_db_utils.py
+++ b/tests/unit/test_db_utils.py
@@ -27,8 +27,6 @@ class TestDatabase(TestCase):
 
         mock_sql.connect.assert_called_with(
             "localhost", "guest", "4321", "users")
-        self.assertIsInstance(
-            results, tuple, "database not created successfully")
 
     @pytest.mark.unit_test
     @mock.patch("gn3.db_utils.SQL_URI",