diff --git a/src/Psalm/Internal/Stubs/CoreGenericFunctions.phpstub b/src/Psalm/Internal/Stubs/CoreGenericFunctions.phpstub
index 251a581c8..a289b6275 100644
--- a/src/Psalm/Internal/Stubs/CoreGenericFunctions.phpstub
+++ b/src/Psalm/Internal/Stubs/CoreGenericFunctions.phpstub
@@ -132,10 +132,11 @@ function array_flip(array $arr)
 
 /**
  * @psalm-template TKey as array-key
+ * @psalm-template TArray as array<TKey, mixed>
  *
- * @param array<TKey, mixed> $arr
+ * @param TArray $arr
  *
- * @return (TKey is empty ? null : ($arr is non-empty-array ? TKey : TKey|null))
+ * @return (TArray is array<empty, empty> ? null : (TArray is non-empty-array ? TKey : TKey|null))
  * @psalm-pure
  */
 function key($arr)
@@ -144,10 +145,11 @@ function key($arr)
 
 /**
  * @psalm-template TKey as array-key
+ * @psalm-template TArray as array<TKey, mixed>
  *
- * @param array<TKey, mixed> $arr
+ * @param TArray $arr
  *
- * @return (TKey is empty ? null : ($arr is non-empty-array ? TKey : TKey|null))
+ * @return (TArray is array<empty, empty> ? null : (TArray is non-empty-array ? TKey : TKey|null))
  * @psalm-pure
  */
 function array_key_first($arr)
@@ -156,10 +158,11 @@ function array_key_first($arr)
 
 /**
  * @psalm-template TKey as array-key
+ * @psalm-template TArray as array<TKey, mixed>
  *
- * @param array<TKey, mixed> $arr
+ * @param TArray $arr
  *
- * @return (TKey is empty ? null : ($arr is non-empty-array ? TKey : TKey|null))
+ * @return (TArray is array<empty, empty> ? null : (TArray is non-empty-array ? TKey : TKey|null))
  * @psalm-pure
  */
 function array_key_last($arr)
diff --git a/tests/ArrayFunctionCallTest.php b/tests/ArrayFunctionCallTest.php
index a2b1771c4..13ad1bc52 100644
--- a/tests/ArrayFunctionCallTest.php
+++ b/tests/ArrayFunctionCallTest.php
@@ -770,6 +770,16 @@ class ArrayFunctionCallTest extends TestCase
                     '$b' => 'string',
                 ],
             ],
+            'keyNonEmptyArray' => [
+                '<?php
+                    /**
+                     * @param non-empty-array $arr
+                     * @return array-key
+                     */
+                    function foo(array $arr) {
+                        return key($arr);
+                    }',
+            ],
             'arrayKeyFirst' => [
                 '<?php
                     /** @return array<string, int> */
diff --git a/tests/TypeReconciliation/EmptyTest.php b/tests/TypeReconciliation/EmptyTest.php
index 5644a9bb4..912fdd6f0 100644
--- a/tests/TypeReconciliation/EmptyTest.php
+++ b/tests/TypeReconciliation/EmptyTest.php
@@ -244,7 +244,7 @@ class EmptyTest extends \Psalm\Tests\TestCase
                         }
                     }',
             ],
-            'SKIPPED-unsetChangesComplicatedArrayEmptiness' => [
+            'unsetChangesComplicatedArrayEmptiness' => [
                 '<?php
                     function contains(array $data, array $needle): bool {
                         if (empty($data) || empty($needle)) {
@@ -253,7 +253,6 @@ class EmptyTest extends \Psalm\Tests\TestCase
                         $stack = [];
 
                         while (!empty($needle)) {
-                            /** @psalm-trace $key */
                             $key = key($needle);
                             $val = $needle[$key];
                             unset($needle[$key]);