1 files changed, 8 insertions, 4 deletions
diff --git a/test/completion/init.lua b/test/completion/init.lua
index 1b2aaf2d..95da8890 100644
--- a/test/completion/init.lua
+++ b/test/completion/init.lua
@@ -110,11 +110,15 @@ function TEST(script)
         end
         assert(result)
         result.complete = nil
-        if expect.include then
-            expect.include = nil
-            assert(include(expect, result))
+        if type(expect) == 'function' then
+            expect(result)
         else
-            assert(eq(expect, result))
+            if expect.include then
+                expect.include = nil
+                assert(include(expect, result))
+            else
+                assert(eq(expect, result))
+            end
         end
         files.remove('')
     end