2 files changed, 11 insertions, 13 deletions
diff --git a/tools/bazeldefs/go.bzl b/tools/bazeldefs/go.bzl
index d388346a5..661c9727e 100644
--- a/tools/bazeldefs/go.bzl
+++ b/tools/bazeldefs/go.bzl
@@ -94,10 +94,10 @@ def go_rule(rule, implementation, **kwargs):
     toolchains = kwargs.get("toolchains", []) + ["@io_bazel_rules_go//go:toolchain"]
     return rule(implementation, attrs = attrs, toolchains = toolchains, **kwargs)
 
-def go_test_library(target):
-    if hasattr(target.attr, "embed") and len(target.attr.embed) > 0:
-        return target.attr.embed[0]
-    return None
+def go_embed_libraries(target):
+    if hasattr(target.attr, "embed"):
+        return target.attr.embed
+    return []
 
 def go_context(ctx, goos = None, goarch = None, std = False):
     """Extracts a standard Go context struct.
diff --git a/tools/nogo/defs.bzl b/tools/nogo/defs.bzl
index b3d297308..161ea972e 100644
--- a/tools/nogo/defs.bzl
+++ b/tools/nogo/defs.bzl
@@ -1,6 +1,6 @@
 """Nogo rules."""
 
-load("//tools/bazeldefs:go.bzl", "go_context", "go_importpath", "go_rule", "go_test_library")
+load("//tools/bazeldefs:go.bzl", "go_context", "go_embed_libraries", "go_importpath", "go_rule")
 
 NogoConfigInfo = provider(
     "information about a nogo configuration",
@@ -200,14 +200,12 @@ def _nogo_aspect_impl(target, ctx):
     # If we're using the "library" attribute, then we need to aggregate the
     # original library sources and dependencies into this target to perform
     # proper type analysis.
-    if ctx.rule.kind == "go_test":
-        library = go_test_library(ctx.rule)
-        if library != None:
-            info = library[NogoInfo]
-            if hasattr(info, "srcs"):
-                srcs = srcs + info.srcs
-            if hasattr(info, "deps"):
-                deps = deps + info.deps
+    for embed in go_embed_libraries(ctx.rule):
+        info = embed[NogoInfo]
+        if hasattr(info, "srcs"):
+            srcs = srcs + info.srcs
+        if hasattr(info, "deps"):
+            deps = deps + info.deps
 
     # Start with all target files and srcs as input.
     inputs = target.files.to_list() + srcs