fix tests

Signed-off-by: Future-Outlier <eric901201@gmail.com>
flyteorg · Jul 25, 2024 · 89d633f · 89d633f
1 parent 05adfd1
commit 89d633f
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 9 deletions.
diff --git a/plugins/flytekit-bigquery/flytekitplugins/bigquery/task.py b/plugins/flytekit-bigquery/flytekitplugins/bigquery/task.py
@@ -38,7 +38,7 @@ def __init__(
         self,
         name: str,
         query_template: str,
-        task_config: Optional[BigQueryConfig],
+        task_config: BigQueryConfig,
         inputs: Optional[Dict[str, Type]] = None,
         output_structured_dataset_type: Optional[Type[StructuredDataset]] = None,
         **kwargs,

diff --git a/plugins/flytekit-snowflake/flytekitplugins/snowflake/task.py b/plugins/flytekit-snowflake/flytekitplugins/snowflake/task.py
@@ -41,7 +41,7 @@ def __init__(
         self,
         name: str,
         query_template: str,
-        task_config: Optional[SnowflakeConfig] = None,
+        task_config: SnowflakeConfig,
         inputs: Optional[Dict[str, Type]] = None,
         output_schema_type: Optional[Type[StructuredDataset]] = None,
         **kwargs,
@@ -64,9 +64,6 @@ def __init__(
                 "results": output_schema_type,
             }
 
-        if task_config is None:
-            task_config = SnowflakeConfig()
-
         super().__init__(
             name=name,
             task_config=task_config,

diff --git a/plugins/flytekit-snowflake/tests/test_agent.py b/plugins/flytekit-snowflake/tests/test_agent.py
@@ -91,17 +91,15 @@ async def test_snowflake_agent(mock_get_private_key):
         schema="dummy_schema",
         warehouse="dummy_warehouse",
         query_id="dummy_id",
+        output=False,
     )
 
     metadata = await agent.create(dummy_template, task_inputs)
     assert metadata == snowflake_metadata
 
     resource = await agent.get(metadata)
     assert resource.phase == TaskExecution.SUCCEEDED
-    assert (
-        resource.outputs.literals["results"].scalar.structured_dataset.uri
-        == "snowflake://dummy_user:dummy_account/dummy_warehouse/dummy_database/dummy_schema/dummy_table"
-    )
+    assert resource.outputs == None
 
     delete_response = await agent.delete(snowflake_metadata)
     assert delete_response is None