[dagster-pipes-test] improve test_message_report_asset_check

Add severity assertion to ensure severity is set correctly. And fix the Java test in PipesTest class to add severity when calling `reportAssetCheck`. Signed-off-by: Thanabodee Charoenpiriyakij <wingyminus@gmail.com>
dagster-io · Dec 25, 2024 · 1a8374d · 1a8374d
1 parent 2102102
commit 1a8374d
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 4 deletions.
diff --git a/...aries/pipes/implementations/java/dagster-pipes-java/src/main/java/pipes/PipesContext.java b/...aries/pipes/implementations/java/dagster-pipes-java/src/main/java/pipes/PipesContext.java
@@ -364,4 +364,4 @@ private String resolveOptionallyPassedAssetKey(
 
         return resultAssetKey;
     }
-}
+}
diff --git a/libraries/pipes/implementations/java/dagster-pipes-java/src/test/java/pipes/MainTest.java b/libraries/pipes/implementations/java/dagster-pipes-java/src/test/java/pipes/MainTest.java
@@ -3,6 +3,7 @@
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import picocli.CommandLine;
+import pipes.data.PipesAssetCheckSeverity;
 import pipes.data.PipesConstants;
 import pipes.loaders.PipesContextLoader;
 import pipes.loaders.PipesDefaultContextLoader;
@@ -183,7 +184,8 @@ public void run() {
                 String checkName = loadParamByWrapperKey("checkName", String.class);
                 boolean passed = loadParamByWrapperKey("passed", Boolean.class);
                 String assetKey = loadParamByWrapperKey("assetKey", String.class);
-                pipesTests.setCheck(checkName, passed, assetKey);
+		PipesAssetCheckSeverity severity = PipesAssetCheckSeverity.valueOf(loadParamByWrapperKey("severity", String.class));
+                pipesTests.setCheck(checkName, passed, assetKey, severity);
             }
 
             if (this.full) {

diff --git a/libraries/pipes/implementations/java/dagster-pipes-java/src/test/java/pipes/PipesTests.java b/libraries/pipes/implementations/java/dagster-pipes-java/src/test/java/pipes/PipesTests.java
@@ -3,6 +3,7 @@
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
+import pipes.data.PipesAssetCheckSeverity;
 import pipes.data.PipesContextData;
 import pipes.data.PipesMetadata;
 import pipes.loaders.*;
@@ -35,6 +36,7 @@ public class PipesTests {
     private String checkName;
     private boolean passed;
     private String checkAssetKey;
+    private PipesAssetCheckSeverity severity;
 
     //Message writer
     private PipesMessageWriter<? extends PipesMessageWriterChannel> pipesMessageWriter;
@@ -70,11 +72,12 @@ void setMaterialization(String dataVersion, String assetKey) {
         this.materializationAssetKey = assetKey;
     }
 
-    void setCheck(String checkName, boolean passed, String assetKey) {
+    void setCheck(String checkName, boolean passed, String assetKey, PipesAssetCheckSeverity severity) {
         this.check = true;
         this.checkName = checkName;
         this.passed = passed;
         this.checkAssetKey = assetKey;
+        this.severity = severity;
     }
 
     void setMessageWriter(PipesMessageWriter<? extends PipesMessageWriterChannel> writer) {
@@ -122,7 +125,7 @@ private void fullTest(PipesContext context) throws DagsterPipesException {
         if (this.check) {
             buildTestMetadata();
             context.reportAssetCheck(
-                this.checkName, this.passed, this.metadata, this.checkAssetKey
+                this.checkName, this.passed, this.severity, this.metadata, this.checkAssetKey
             );
         }
     }

diff --git a/libraries/pipes/tests/dagster-pipes-tests/src/dagster_pipes_tests/suite.py b/libraries/pipes/tests/dagster-pipes-tests/src/dagster_pipes_tests/suite.py
@@ -11,6 +11,7 @@
 import pytest
 from dagster import (
     AssetCheckResult,
+    AssetCheckSeverity,
     AssetCheckSpec,
     AssetExecutionContext,
     AssetKey,
@@ -613,6 +614,8 @@ def my_asset(
             if check_result.metadata is not None:
                 assert_known_metadata(check_result.metadata)  # type: ignore
 
+            assert check_result.severity == AssetCheckSeverity(severity)
+
             yield from results
 
         result = materialize(
-Original file line number
+Diff line change
@@ Expand Up / @@ -364,4 +364,4 @@ private String resolveOptionallyPassedAssetKey( @@
             return resultAssetKey;
         }
-    }
+    }