From 192c992e219e1f22b97f9f24e683c9597613f20f Mon Sep 17 00:00:00 2001
From: Pere Miquel Brull <peremiquelbrull@gmail.com>
Date: Fri, 27 Jan 2023 12:55:57 +0100
Subject: [PATCH] Fix e2e tests when computing rows (#9958)

* Add vertica to the list

* test profiler

* Fix tests

* Fix tests
---
 .github/workflows/py-cli-e2e-tests.yml        |  2 +-
 .../cli_e2e/database/snowflake/snowflake.yaml |  2 +-
 ingestion/tests/cli_e2e/test_cli_db_base.py   |  4 ++-
 ingestion/tests/cli_e2e/test_cli_mssql.py     |  5 ++--
 ingestion/tests/cli_e2e/test_cli_snowflake.py | 26 ++++++++++++++++++-
 5 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/py-cli-e2e-tests.yml b/.github/workflows/py-cli-e2e-tests.yml
index ecbffbea7d9..e2243be7e11 100644
--- a/.github/workflows/py-cli-e2e-tests.yml
+++ b/.github/workflows/py-cli-e2e-tests.yml
@@ -21,7 +21,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        e2e-test: ['python', 'mysql', 'bigquery', 'snowflake', 'dbt_redshift', 'mssql']
+        e2e-test: ['python', 'mysql', 'bigquery', 'snowflake', 'dbt_redshift', 'mssql', 'vertica']
     environment: test
 
     steps:
diff --git a/ingestion/tests/cli_e2e/database/snowflake/snowflake.yaml b/ingestion/tests/cli_e2e/database/snowflake/snowflake.yaml
index 32e96f8cf3d..40c8b75a9c5 100644
--- a/ingestion/tests/cli_e2e/database/snowflake/snowflake.yaml
+++ b/ingestion/tests/cli_e2e/database/snowflake/snowflake.yaml
@@ -1,6 +1,6 @@
 source:
   type: snowflake
-  serviceName: local_snowflake
+  serviceName: e2e_snowflake
   serviceConnection:
     config:
       username: $E2E_SNOWFLAKE_USERNAME
diff --git a/ingestion/tests/cli_e2e/test_cli_db_base.py b/ingestion/tests/cli_e2e/test_cli_db_base.py
index 7631a18edb7..6b56504934f 100644
--- a/ingestion/tests/cli_e2e/test_cli_db_base.py
+++ b/ingestion/tests/cli_e2e/test_cli_db_base.py
@@ -349,7 +349,9 @@ class CliDBBase(TestCase):
                     "config": {
                         "type": "Profiler",
                         "generateSampleData": True,
-                        "profileSample": 1,
+                        "profileSample": extra_args.get("profileSample", 1)
+                        if extra_args
+                        else 1,
                     }
                 }
                 config_yaml["processor"] = {"type": "orm-profiler", "config": {}}
diff --git a/ingestion/tests/cli_e2e/test_cli_mssql.py b/ingestion/tests/cli_e2e/test_cli_mssql.py
index a3c4ace0a18..f54ea77a1fb 100644
--- a/ingestion/tests/cli_e2e/test_cli_mssql.py
+++ b/ingestion/tests/cli_e2e/test_cli_mssql.py
@@ -47,8 +47,7 @@ class MSSQLCliTest(CliCommonDB.TestSuite, SQACommonMethods):
         (3,'Steve Rogers', '1988-07-04'),
         (4,'Natasha Romanoff', '1997-12-03'),
         (5,'Wanda Maximoff', '1998-02-10'),
-        (6,'Diana Prince', '1976-03-17')
-        ;
+        (6,'Diana Prince', '1976-03-17');
     """
     ]
 
@@ -123,7 +122,7 @@ class MSSQLCliTest(CliCommonDB.TestSuite, SQACommonMethods):
         return 1
 
     def inserted_rows_count(self) -> int:
-        return len(self.insert_data_queries)
+        return 6
 
     @staticmethod
     def fqn_created_table() -> str:
diff --git a/ingestion/tests/cli_e2e/test_cli_snowflake.py b/ingestion/tests/cli_e2e/test_cli_snowflake.py
index c0de922687c..02f13b81033 100644
--- a/ingestion/tests/cli_e2e/test_cli_snowflake.py
+++ b/ingestion/tests/cli_e2e/test_cli_snowflake.py
@@ -14,9 +14,12 @@ Test Snowflake connector with CLI
 """
 from typing import List
 
+import pytest
+
 from metadata.ingestion.api.sink import SinkStatus
 from metadata.ingestion.api.source import SourceStatus
 
+from .test_cli_db_base import E2EType
 from .test_cli_db_base_common import CliCommonDB
 
 
@@ -80,6 +83,27 @@ class SnowflakeCliTest(CliCommonDB.TestSuite):
             connection.execute(self.drop_table_query)
             connection.close()
 
+    @pytest.mark.order(2)
+    def test_create_table_with_profiler(self) -> None:
+        # delete table in case it exists
+        self.delete_table_and_view()
+        # create a table and a view
+        self.create_table_and_view()
+        # build config file for ingest
+        self.build_config_file()
+        # run ingest with new tables
+        self.run_command()
+        # build config file for profiler
+        self.build_config_file(
+            E2EType.PROFILER,
+            # Otherwise the sampling here does not pick up rows
+            extra_args={"profileSample": 100},
+        )
+        # run profiler with new tables
+        result = self.run_command("profile")
+        sink_status, source_status = self.retrieve_statuses(result)
+        self.assert_for_table_with_profiler(source_status, sink_status)
+
     @staticmethod
     def expected_tables() -> int:
         return 7
@@ -89,7 +113,7 @@ class SnowflakeCliTest(CliCommonDB.TestSuite):
 
     @staticmethod
     def fqn_created_table() -> str:
-        return "local_snowflake.E2E_DB.E2E_TEST.PERSONS"
+        return "e2e_snowflake.E2E_DB.E2E_TEST.PERSONS"
 
     @staticmethod
     def get_includes_schemas() -> List[str]: