fix(nix): refactor nixosTest injections for better performance (#8378)

2 files changed, 31 insertions, 13 deletions

diff --git a/runtime/queries/nix/injections.scm b/runtime/queries/nix/injections.scm
index 7d0ffe8e5..718a9c1c5 100644
--- a/runtime/queries/nix/injections.scm
+++ b/runtime/queries/nix/injections.scm
@@ -251,17 +251,35 @@
   (#lua-match? @_func "^%a*%.*writeRust%a*$")
   (#set! injection.combined))
 
+; (runTest) testScript
+(apply_expression
+  function: (_) @_func
+  argument: (_
+    (_)*
+    (_
+      (binding
+        attrpath: (attrpath) @_func_name
+        expression: (_
+          (string_fragment) @injection.content
+          (#set! injection.language "python")))
+      (#eq? @_func_name "testScript")
+      (#lua-match? @_func "^.*%.*runTest$")
+      (#set! injection.combined))))
+
 ; (nixosTest) testScript
-((binding
-  attrpath: (attrpath) @_attr_name
-  (#eq? @_attr_name "nodes"))
-  (binding
-    attrpath: (attrpath) @_func_name
-    (#eq? @_func_name "testScript")
+(apply_expression
+  function: (_) @_func
+  argument: (_
+    (_)*
     (_
-      (string_fragment) @injection.content
-      (#set! injection.language "python")))
-  (#set! injection.combined))
+      (binding
+        attrpath: (attrpath) @_func_name
+        expression: (_
+          (string_fragment) @injection.content
+          (#set! injection.language "python")))
+      (#eq? @_func_name "testScript")
+      (#lua-match? @_func "^.*%.*nixosTest$")
+      (#set! injection.combined))))
 
 ; home-manager Neovim plugin config
 (attrset_expression
diff --git a/tests/query/injections/nix/test-nix-injections.nix b/tests/query/injections/nix/test-nix-injections.nix
index b358326ab..7505398e6 100644
--- a/tests/query/injections/nix/test-nix-injections.nix
+++ b/tests/query/injections/nix/test-nix-injections.nix
@@ -1,7 +1,7 @@
 # NOTE: This file is for testing queries. If evaluated with nix, this will probably error.
-{ pkgs, stdenv }: let
+{ pkgs, stdenv, lib }: let
   match = builtins.match;
-in { 
+in {
   regex = builtins.match "^.*[{](.+)[}].*$" "blahblah {something} blahblah" [ "something" ];
 
   regex2 = match "^.*[{](.+)[}].*$" "blahblah {something} blahblah" [ "something" ];
@@ -12,7 +12,7 @@ in {
       echo "bar" > $out/foo.txt
     '';
   };
-  
+
   drv2 = pkgs.writeShellApplication {
     name = "shellApp";
     buildInputs = [];
@@ -62,7 +62,7 @@ in {
     }
   '';
 
-  drv12 = {
+  drv12 = lib.nixos.runTest {
     nodes = null;
     testScript = ''
       print("Hello world!");