diff --git a/src/__tests__/lib/rules/prefer-in-document.js b/src/__tests__/lib/rules/prefer-in-document.js
index 917a619..75aac07 100644
--- a/src/__tests__/lib/rules/prefer-in-document.js
+++ b/src/__tests__/lib/rules/prefer-in-document.js
@@ -35,6 +35,8 @@ const valid = [
   ]),
   `expect(screen.notAQuery('foo-bar')).toHaveLength(1)`,
   `expect(screen.getByText('foo-bar')).toHaveLength(2)`,
+  `expect(undefinedVariable).toBeDefined()`,
+  `expect([1,2,3]).toHaveLength(1);`,
 ];
 const invalid = [
   // Invalid cases that applies to all variants
@@ -51,6 +53,14 @@ const invalid = [
       `expect(wrapper.${q}('foo')).toHaveLength(1)`,
       `expect(wrapper.${q}('foo')).toBeInTheDocument()`
     ),
+    invalidCase(
+      `const element = screen.${q}('foo'); expect(element).toHaveLength(1);`,
+      `const element = screen.${q}('foo'); expect(element).toBeInTheDocument();`
+    ),
+    invalidCase(
+      `const element = ${q}('foo'); expect(element).toBeNull();`,
+      `const element = ${q}('foo'); expect(element).not.toBeInTheDocument();`
+    ),
   ]),
   // Invalid cases that applies to queryBy* and queryAllBy*
   ...queriesByVariant.query.map((q) => [
diff --git a/src/rules/prefer-in-document.js b/src/rules/prefer-in-document.js
index fd06149..ec7a0b5 100644
--- a/src/rules/prefer-in-document.js
+++ b/src/rules/prefer-in-document.js
@@ -29,15 +29,33 @@ function isAntonymMatcher(matcherNode, matcherArguments) {
 
 function check(
   context,
-  { queryNode, matcherNode, matcherArguments, negatedMatcher }
+  { subject, matcherNode, matcherArguments, negatedMatcher }
 ) {
-  const query = queryNode.name || queryNode.property.name;
+  let query;
+
+  if (subject.type === "Identifier") {
+    // Backtrack the variable to see if it was populated by a query.
+    try {
+      const variable = context
+        .getScope()
+        .variables.find((v) => v.name === subject.name).defs[0].node.init;
+      query = variable.callee.name || variable.callee.property.name;
+    } catch (error) {
+      return;
+    }
+  } else if (subject.type !== "CallExpression") {
+    return;
+  }
 
   // toHaveLength() is only invalid with 0 or 1
   if (matcherNode.name === "toHaveLength" && matcherArguments[0].value > 1) {
     return;
   }
 
+  if (!query) {
+    query = subject.callee.name || subject.callee.property.name;
+  }
+
   if (queries.includes(query)) {
     context.report({
       node: matcherNode,
@@ -82,13 +100,13 @@ export const create = (context) => {
     [`CallExpression[callee.object.object.callee.name='expect'][callee.object.property.name='not'][callee.property.name=${alternativeMatchers}]`](
       node
     ) {
-      const queryNode = node.callee.object.object.arguments[0].callee;
+      const subject = node.callee.object.object.arguments[0];
       const matcherNode = node.callee.property;
       const matcherArguments = node.arguments;
 
       check(context, {
         negatedMatcher: true,
-        queryNode,
+        subject,
         matcherNode,
         matcherArguments,
       });
@@ -98,13 +116,13 @@ export const create = (context) => {
     [`CallExpression[callee.object.callee.name='expect'][callee.property.name=${alternativeMatchers}]`](
       node
     ) {
-      const queryNode = node.callee.object.arguments[0].callee;
+      const subject = node.callee.object.arguments[0];
       const matcherNode = node.callee.property;
       const matcherArguments = node.arguments;
 
       check(context, {
         negatedMatcher: false,
-        queryNode,
+        subject,
         matcherNode,
         matcherArguments,
       });