diff --git a/mlir/test/Dialect/SPIRV/IR/target-env.mlir b/mlir/test/Dialect/SPIRV/IR/target-env.mlir
--- a/mlir/test/Dialect/SPIRV/IR/target-env.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/target-env.mlir
@@ -1,4 +1,5 @@
-// RUN: mlir-opt -mlir-disable-threading -test-spirv-target-env %s | FileCheck %s
+// RUN: mlir-opt --split-input-file --verify-diagnostics --mlir-disable-threading \
+// RUN:   --test-spirv-target-env %s | FileCheck %s
 
 // Note: The following tests check that a spirv.target_env can properly control
 // the conversion target and filter unavailable ops during the conversion.
@@ -32,6 +33,14 @@
 // Vulkan memory model is available via extension SPV_KHR_vulkan_memory_model,
 // which extensions are incorporated into SPIR-V 1.5.
 
+
+// expected-error @+1 {{missing 'spirv.target_env' attribute}}
+func.func @main() {
+  return
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // MaxVersion
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/SPIRV/TestAvailability.cpp b/mlir/test/lib/Dialect/SPIRV/TestAvailability.cpp
--- a/mlir/test/lib/Dialect/SPIRV/TestAvailability.cpp
+++ b/mlir/test/lib/Dialect/SPIRV/TestAvailability.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
 #include "mlir/Pass/Pass.h"
@@ -210,9 +211,8 @@
   MLIRContext *context = &getContext();
   func::FuncOp fn = getOperation();
 
-  auto targetEnv = fn.getOperation()
-                       ->getAttr(spirv::getTargetEnvAttrName())
-                       .cast<spirv::TargetEnvAttr>();
+  auto targetEnv = dyn_cast_or_null<spirv::TargetEnvAttr>(
+      fn.getOperation()->getAttr(spirv::getTargetEnvAttrName()));
   if (!targetEnv) {
     fn.emitError("missing 'spirv.target_env' attribute");
     return signalPassFailure();