From bd7bb0201c28e95e99eb0f446e69528569ed1d76 Mon Sep 17 00:00:00 2001
From: vqmarkman <104043363+vqmarkman@users.noreply.github.com>
Date: Wed, 6 Sep 2023 11:44:03 -0700
Subject: [PATCH] Issue #152: probes unit tests. (#253)

- renamed one of the tests to reflect what it actually does.
	    - when  #241 is fixed, tests will need to be adjusted as well.
---
 test/common_utils.py            |  12 ++
 test/conftest.py                |   8 +-
 test/unit/test_probe_procs.py   | 237 ++++++++++++++++++--------------
 test/unit/test_probe_upgrade.py | 111 +++++++++++++++
 4 files changed, 267 insertions(+), 101 deletions(-)
 create mode 100644 test/unit/test_probe_upgrade.py

diff --git a/test/common_utils.py b/test/common_utils.py
index 2e3e1866..996d045d 100644
--- a/test/common_utils.py
+++ b/test/common_utils.py
@@ -38,3 +38,15 @@ def delete_list_of_labels(conn, sql):
             assert "done" in str(
                 run_proc(conn, delete_label_statement)
             ), "Stored procedure output does not match expected result!"
+
+
+def delete_list_of_probes(conn, sql):
+    print(f"[INFO] SQL in delete function: {sql}")
+
+    with conn() as cnx:
+        cur = cnx.cursor()
+        for name in cur.execute(sql).fetchall():
+            delete_probe_statement = f"call ADMIN.DELETE_PROBE('{name[0]}');"
+            assert "done" in str(
+                run_proc(conn, delete_probe_statement)
+            ), "Stored procedure output does not match expected result!"
diff --git a/test/conftest.py b/test/conftest.py
index 83ed771a..b6c19093 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -2,7 +2,7 @@
 import datetime
 import pytest
 from contextlib import contextmanager
-from common_utils import delete_list_of_labels
+from common_utils import delete_list_of_labels, delete_list_of_probes
 
 sys.path.append("../deploy")
 import helpers  # noqa E402
@@ -62,3 +62,9 @@ def timestamp_string(conn):
 
     # call a function that deletes all the labels that were created in the session
     delete_list_of_labels(conn, sql)
+
+    sql = f"select name from INTERNAL.PROBES where name like '%{ts}%'"
+    print(f"[INFO] SQL stmt to find all the probes: {sql}")
+
+    # call a function that deletes all the labels that were created in the session
+    delete_list_of_probes(conn, sql)
diff --git a/test/unit/test_probe_procs.py b/test/unit/test_probe_procs.py
index a558e9ae..ff8c6170 100644
--- a/test/unit/test_probe_procs.py
+++ b/test/unit/test_probe_procs.py
@@ -1,111 +1,148 @@
 from __future__ import annotations
 
+import pytest
+from common_utils import generate_unique_name
 from common_utils import run_proc
 from common_utils import row_count
-from common_utils import run_sql
-import time
 
 
-def test_initialize_then_migrate_probes(conn, timestamp_string):
-    # step 1: clean up the probes table and predefined_probes table
-    sql = "truncate table internal.probes"
-    assert "successfully" in str(
-        run_sql(conn, sql)
-    ), "SQL output does not match expected result!"
+def test_smoke_create_drop_probe(conn, timestamp_string):
+    probe = generate_unique_name("probe", timestamp_string)
+    sql = f"CALL ADMIN.CREATE_PROBE('{probe}', 'rows_produced > 100', True :: BOOLEAN, 'SLACK', 'jinfeng@sundeck.io', 'SLACK', False :: BOOLEAN);"
+
+    # create_probe returns NULL in case of successful probe creation
+    assert run_proc(conn, sql) is None, "Stored procedure did not return NULL value!"
+
+    # make sure it was created with correct properties
+    sql = f"""select count(*) from INTERNAL.probes where
+                 name = '{probe}' and
+                 condition = 'rows_produced > 100' and
+                 notify_writer and
+                 notify_writer_method = 'SLACK' and
+                 notify_other = 'jinfeng@sundeck.io' and
+                 notify_other_method = 'SLACK' and
+                 probe_created_at is null and
+                 probe_modified_at is not null and
+                 not cancel and
+                 enabled is null
+        """
+    assert row_count(conn, sql) == 1, "Probe was not found!"
+
+    # drop probe
+    sql = f"call ADMIN.DELETE_PROBE('{probe}');"
+    assert run_proc(conn, sql) == "done", "Stored procedure did not return NULL value!"
+
+
+# Test that validates that we get correct error on attempt to create probe with existing name
+def test_create_probe_with_existing_name(conn, timestamp_string):
+
+    probe = generate_unique_name("probe", timestamp_string)
+    sql = f"CALL ADMIN.CREATE_PROBE('{probe}', 'rows_produced > 100', True :: BOOLEAN, 'SLACK', 'jinfeng@sundeck.io', 'SLACK', False :: BOOLEAN);"
+    assert run_proc(conn, sql) is None, "Stored procedure did not return NULL value!"
+
+    assert (
+        run_proc(conn, sql)
+        == "A probe with this name already exists. Please choose a distinct name."
+    ), "Stored procedure output does not match expected result!"
 
-    sql = "truncate table internal.predefined_probes"
-    assert "successfully" in str(
-        run_sql(conn, sql)
-    ), "SQL output does not match expected result!"
 
-    # step 2: clean up the flag in internal.config
-    sql = "delete from internal.config where KEY = 'PROBES_INITED'"
-    run_sql(conn, sql)
+def test_smoke_update_probe(conn, timestamp_string):
+    probe = generate_unique_name("probe", timestamp_string)
+    sql = f"CALL ADMIN.CREATE_PROBE('{probe}', 'compilation_time  > 50000', True :: BOOLEAN, 'EMAIL', 'doron@sundeck.io', 'EMAIL', False :: BOOLEAN);"
+
+    # create_probe returns NULL in case of successful probe creation
+    assert run_proc(conn, sql) is None, "Stored procedure did not return NULL value!"
+
+    # make sure it was created with correct properties
+    sql = f"""select count(*) from INTERNAL.probes where
+                 name = '{probe}' and
+                 condition = 'compilation_time  > 50000' and
+                 notify_writer and
+                 notify_writer_method = 'EMAIL' and
+                 notify_other = 'doron@sundeck.io' and
+                 notify_other_method = 'EMAIL' and
+                 probe_created_at is null and
+                 probe_modified_at is not null and
+                 not cancel and
+                 enabled is null
+        """
+    assert row_count(conn, sql) == 1, "Label was not found!"
+
+    # update probe
+    sql = f"CALL ADMIN.UPDATE_PROBE('{probe}', '{probe}', 'rows_produced = 1000', True :: BOOLEAN, 'SLACK', 'doron@sundeck.io', 'SLACK', False :: BOOLEAN);"
+    assert run_proc(conn, sql) is None, "Stored procedure did not return NULL value!"
 
-    # step 3: populate predefined_probes table
-    sql = "CALL INTERNAL.POPULATE_PREDEFINED_PROBES();"
+    # validate that probe was updated correctly
+    sql = f"""select count(*) from INTERNAL.probes where
+                 name = '{probe}' and
+                 condition = 'rows_produced = 1000' and
+                 notify_writer and
+                 notify_writer_method = 'SLACK' and
+                 notify_other = 'doron@sundeck.io' and
+                 notify_other_method = 'SLACK' and
+                 probe_created_at is null and
+                 probe_modified_at is not null and
+                 not cancel and
+                 enabled is null
+        """
+    assert row_count(conn, sql) == 1, "Probe was not found!"
+
+
+# Test that validates that we can create/drop probe with empty string for name
+# Legal in Snowflake
+def test_create_probe_with_empty_string_name(conn, timestamp_string):
+
+    sql = "CALL ADMIN.CREATE_PROBE('', 'compilation_time  > 50000', True :: BOOLEAN, 'EMAIL', 'doron@sundeck.io', 'EMAIL', False :: BOOLEAN);"
     assert run_proc(conn, sql) is None, "Stored procedure did not return NULL value!"
 
-    sql = "select count(*) from internal.PREDEFINED_PROBES"
-    output = row_count(conn, sql)
-    assert output > 0, "SQL output " + str(output) + " does not match expected result!"
-
-    # step 4: call internal.initialize_probes()
-    sql = "call INTERNAL.INITIALIZE_PROBES()"
-    output = str(run_sql(conn, sql))
-    assert "True" in output, "SQL output" + output + " does not match expected result!"
-
-    # step 5: verify rows in probes table
-    sql = "select count(*) from internal.PROBES"
-    output = row_count(conn, sql)
-    assert output > 0, "SQL output " + str(output) + " does not match expected result!"
-
-    # step 6: verify flag in internal.config
-    sql = "call internal.get_config('PROBES_INITED')"
-    output = str(run_sql(conn, sql))
-    assert "True" in output, "SQL output" + output + " does not match expected result!"
-
-    # step 7: call internal.initialize_probes() again
-    sql = "call INTERNAL.INITIALIZE_PROBES()"
-    output = str(run_sql(conn, sql))
-    assert "False" in output, "SQL output" + output + " does not match expected result!"
-
-    # step 8: verify rows in probes table
-    sql = "select count(*) from internal.PROBES"
-    output = row_count(conn, sql)
-    assert output > 0, "SQL output " + str(output) + " does not match expected result!"
-
-    # step 9: sleep 5 seconds
-    time.sleep(5)
-
-    # step 10: call internal.migrate_predefined_probes()
-    sql = "call INTERNAL.MIGRATE_PREDEFINED_PROBES(5)"
-    output = str(run_sql(conn, sql))
-    assert "True" in output, "SQL output" + output + " does not match expected result!"
-
-    # step 11: insert a new predefined probe to PREDEFIEND_PROBES
-    sql = "INSERT INTO INTERNAL.PREDEFINED_PROBES (name, condition, PROBE_CREATED_AT, PROBE_MODIFIED_AT) values ('NEW PREDEFINED PROBE', 'bytes_scanned > 10000', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)"
-    run_sql(conn, sql)
-
-    # step 12: call internal.migrate_predefined_probes(). Migration should return true, since we have a new predefined probe.
-    sql = "call INTERNAL.MIGRATE_PREDEFINED_PROBES(5)"
-    output = str(run_sql(conn, sql))
-    assert "True" in output, "SQL output" + output + " does not match expected result!"
-
-    # step 13: verify probes table has the new added predefinend probe "NEW PREDEFINED PROBE"
-    sql = "select count(*) from internal.PROBES where name = 'NEW PREDEFINED PROBE'"
-    rowcount = row_count(conn, sql)
-    assert rowcount == 1, (
-        "SQL output " + str(rowcount) + " does not match expected result!"
-    )
-
-    # step 14: update the condition of 'NEW PREDEFINED PROBE'
-    sql = "UPDATE INTERNAL.PREDEFINED_PROBES SET CONDITION = 'bytes_scanned > 20000 ' where NAME = 'NEW PREDEFINED PROBE'"
-    run_sql(conn, sql)
-
-    # step 15: call internal.migrate_predefined_probes(). Migration should return true, since we modify condition of one old predefined probes.
-    time.sleep(5)
-
-    sql = "call INTERNAL.MIGRATE_PREDEFINED_probes(5)"
-    output = str(run_sql(conn, sql))
-    assert "True" in output, "SQL output" + output + " does not match expected result!"
-
-    # step 16: insert a row into user's PROBES
-    sql = "INSERT INTO INTERNAL.PROBES (name, condition, PROBE_CREATED_AT, PROBE_MODIFIED_AT) values ('test', 'bytes_scanned > 100', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)"
-    run_sql(conn, sql)
-
-    ## step 17: MIGRATE_PREDEFINED_PROBES should return False, because user adds one PROBE
-    sql = "call INTERNAL.MIGRATE_PREDEFINED_PROBES(5)"
-    output = str(run_sql(conn, sql))
-    assert "False" in output, "SQL output" + output + " does not match expected result!"
-
-    # step 18: clean up the probes table and predefined_probes table
-    sql = "truncate table internal.probes"
-    assert "successfully" in str(
-        run_sql(conn, sql)
-    ), "SQL output does not match expected result!"
-
-    sql = "truncate table internal.predefined_probes"
-    assert "successfully" in str(
-        run_sql(conn, sql)
-    ), "SQL output does not match expected result!"
+    sql = "call ADMIN.DELETE_PROBE('');"
+    assert "done" in str(
+        run_proc(conn, sql)
+    ), "Stored procedure output does not match expected result!"
+
+
+# List of test cases with statements and expected error messages
+test_cases = [
+    (
+        "CALL ADMIN.CREATE_PROBE(NULL, 'rows_produced > 100', True :: BOOLEAN, 'SLACK', 'jinfeng@sundeck.io', 'SLACK', False :: BOOLEAN);",
+        "Name must not be null.",
+    ),
+    (
+        "CALL ADMIN.CREATE_PROBE('{probe}', NULL, True :: BOOLEAN, 'SLACK', 'jinfeng@sundeck.io', 'SLACK', False :: BOOLEAN);",
+        "Condition must not be null.",
+    ),
+    (
+        "CALL ADMIN.CREATE_PROBE('{probe}', 'x=y and z is not null', True :: BOOLEAN, 'SLACK', 'jinfeng@sundeck.io', 'SLACK', False :: BOOLEAN);",
+        "Invalid condition SQL. Please check your syntax.",
+    ),
+    (
+        "CALL ADMIN.CREATE_PROBE('{probe}', 'blah blah and ', True :: BOOLEAN, 'SLACK', 'jinfeng@sundeck.io', 'SLACK', False :: BOOLEAN);",
+        "Invalid condition SQL. Please check your syntax.",
+    ),
+    (
+        "CALL ADMIN.UPDATE_PROBE('{probe}', NULL, 'compilation_time > 3000', True :: BOOLEAN, 'SLACK', 'doron@sundeck.io', 'SLACK', False :: BOOLEAN);",
+        "Name must not be null.",
+    ),
+    (
+        "CALL ADMIN.UPDATE_PROBE('{probe}', 'new_probe_name', NULL, True :: BOOLEAN, 'SLACK', 'doron@sundeck.io', 'SLACK', False :: BOOLEAN);",
+        "Condition must not be null.",
+    ),
+    (
+        "CALL ADMIN.UPDATE_PROBE('{probe}', 'new_probe_name', 'x=y and z is not null', True :: BOOLEAN, 'SLACK', 'doron@sundeck.io', 'SLACK', False :: BOOLEAN);",
+        "Invalid condition SQL. Please check your syntax.",
+    ),
+    (
+        "CALL ADMIN.UPDATE_PROBE('{probe}', 'new_probe_name', 'x=y and z is not null', True :: BOOLEAN, 'SLACK', 'doron@sundeck.io', 'SLACK', False :: BOOLEAN);",
+        "Invalid condition SQL. Please check your syntax.",
+    ),
+]
+
+# Test that validates that correct error message was returned
+@pytest.mark.parametrize("statement, expected_error", test_cases)
+def test_error_message(conn, timestamp_string, statement, expected_error):
+
+    probe = generate_unique_name("probe", timestamp_string)
+    sql = statement.format(probe=probe)
+    assert expected_error in str(
+        run_proc(conn, sql)
+    ), "Stored procedure output does not match expected result!"
diff --git a/test/unit/test_probe_upgrade.py b/test/unit/test_probe_upgrade.py
new file mode 100644
index 00000000..a558e9ae
--- /dev/null
+++ b/test/unit/test_probe_upgrade.py
@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+from common_utils import run_proc
+from common_utils import row_count
+from common_utils import run_sql
+import time
+
+
+def test_initialize_then_migrate_probes(conn, timestamp_string):
+    # step 1: clean up the probes table and predefined_probes table
+    sql = "truncate table internal.probes"
+    assert "successfully" in str(
+        run_sql(conn, sql)
+    ), "SQL output does not match expected result!"
+
+    sql = "truncate table internal.predefined_probes"
+    assert "successfully" in str(
+        run_sql(conn, sql)
+    ), "SQL output does not match expected result!"
+
+    # step 2: clean up the flag in internal.config
+    sql = "delete from internal.config where KEY = 'PROBES_INITED'"
+    run_sql(conn, sql)
+
+    # step 3: populate predefined_probes table
+    sql = "CALL INTERNAL.POPULATE_PREDEFINED_PROBES();"
+    assert run_proc(conn, sql) is None, "Stored procedure did not return NULL value!"
+
+    sql = "select count(*) from internal.PREDEFINED_PROBES"
+    output = row_count(conn, sql)
+    assert output > 0, "SQL output " + str(output) + " does not match expected result!"
+
+    # step 4: call internal.initialize_probes()
+    sql = "call INTERNAL.INITIALIZE_PROBES()"
+    output = str(run_sql(conn, sql))
+    assert "True" in output, "SQL output" + output + " does not match expected result!"
+
+    # step 5: verify rows in probes table
+    sql = "select count(*) from internal.PROBES"
+    output = row_count(conn, sql)
+    assert output > 0, "SQL output " + str(output) + " does not match expected result!"
+
+    # step 6: verify flag in internal.config
+    sql = "call internal.get_config('PROBES_INITED')"
+    output = str(run_sql(conn, sql))
+    assert "True" in output, "SQL output" + output + " does not match expected result!"
+
+    # step 7: call internal.initialize_probes() again
+    sql = "call INTERNAL.INITIALIZE_PROBES()"
+    output = str(run_sql(conn, sql))
+    assert "False" in output, "SQL output" + output + " does not match expected result!"
+
+    # step 8: verify rows in probes table
+    sql = "select count(*) from internal.PROBES"
+    output = row_count(conn, sql)
+    assert output > 0, "SQL output " + str(output) + " does not match expected result!"
+
+    # step 9: sleep 5 seconds
+    time.sleep(5)
+
+    # step 10: call internal.migrate_predefined_probes()
+    sql = "call INTERNAL.MIGRATE_PREDEFINED_PROBES(5)"
+    output = str(run_sql(conn, sql))
+    assert "True" in output, "SQL output" + output + " does not match expected result!"
+
+    # step 11: insert a new predefined probe to PREDEFIEND_PROBES
+    sql = "INSERT INTO INTERNAL.PREDEFINED_PROBES (name, condition, PROBE_CREATED_AT, PROBE_MODIFIED_AT) values ('NEW PREDEFINED PROBE', 'bytes_scanned > 10000', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)"
+    run_sql(conn, sql)
+
+    # step 12: call internal.migrate_predefined_probes(). Migration should return true, since we have a new predefined probe.
+    sql = "call INTERNAL.MIGRATE_PREDEFINED_PROBES(5)"
+    output = str(run_sql(conn, sql))
+    assert "True" in output, "SQL output" + output + " does not match expected result!"
+
+    # step 13: verify probes table has the new added predefinend probe "NEW PREDEFINED PROBE"
+    sql = "select count(*) from internal.PROBES where name = 'NEW PREDEFINED PROBE'"
+    rowcount = row_count(conn, sql)
+    assert rowcount == 1, (
+        "SQL output " + str(rowcount) + " does not match expected result!"
+    )
+
+    # step 14: update the condition of 'NEW PREDEFINED PROBE'
+    sql = "UPDATE INTERNAL.PREDEFINED_PROBES SET CONDITION = 'bytes_scanned > 20000 ' where NAME = 'NEW PREDEFINED PROBE'"
+    run_sql(conn, sql)
+
+    # step 15: call internal.migrate_predefined_probes(). Migration should return true, since we modify condition of one old predefined probes.
+    time.sleep(5)
+
+    sql = "call INTERNAL.MIGRATE_PREDEFINED_probes(5)"
+    output = str(run_sql(conn, sql))
+    assert "True" in output, "SQL output" + output + " does not match expected result!"
+
+    # step 16: insert a row into user's PROBES
+    sql = "INSERT INTO INTERNAL.PROBES (name, condition, PROBE_CREATED_AT, PROBE_MODIFIED_AT) values ('test', 'bytes_scanned > 100', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)"
+    run_sql(conn, sql)
+
+    ## step 17: MIGRATE_PREDEFINED_PROBES should return False, because user adds one PROBE
+    sql = "call INTERNAL.MIGRATE_PREDEFINED_PROBES(5)"
+    output = str(run_sql(conn, sql))
+    assert "False" in output, "SQL output" + output + " does not match expected result!"
+
+    # step 18: clean up the probes table and predefined_probes table
+    sql = "truncate table internal.probes"
+    assert "successfully" in str(
+        run_sql(conn, sql)
+    ), "SQL output does not match expected result!"
+
+    sql = "truncate table internal.predefined_probes"
+    assert "successfully" in str(
+        run_sql(conn, sql)
+    ), "SQL output does not match expected result!"