diff --git a/lib/gitlab/ci/pipeline/expression/lexer.rb b/lib/gitlab/ci/pipeline/expression/lexer.rb
index cabb91d0dd7..e1c68b7c3c2 100644
--- a/lib/gitlab/ci/pipeline/expression/lexer.rb
+++ b/lib/gitlab/ci/pipeline/expression/lexer.rb
@@ -16,12 +16,13 @@ module Gitlab
 
           MAX_TOKENS = 100
 
-          def initialize(statement)
+          def initialize(statement, max_tokens: MAX_TOKENS)
             @scanner = StringScanner.new(statement)
+            @max_tokens = max_tokens
           end
 
-          def tokens(max: MAX_TOKENS)
-            strong_memoize(:tokens) { tokenize(max) }
+          def tokens
+            strong_memoize(:tokens) { tokenize }
           end
 
           def lexemes
@@ -30,10 +31,10 @@ module Gitlab
 
           private
 
-          def tokenize(max_tokens)
+          def tokenize
             tokens = []
 
-            max_tokens.times do
+            @max_tokens.times do
               @scanner.skip(/\s+/) # ignore whitespace
 
               return tokens if @scanner.eos?
diff --git a/spec/lib/gitlab/ci/pipeline/expression/lexer_spec.rb b/spec/lib/gitlab/ci/pipeline/expression/lexer_spec.rb
index 65f90b0caf3..230ceeb07f8 100644
--- a/spec/lib/gitlab/ci/pipeline/expression/lexer_spec.rb
+++ b/spec/lib/gitlab/ci/pipeline/expression/lexer_spec.rb
@@ -46,9 +46,9 @@ describe Gitlab::Ci::Pipeline::Expression::Lexer do
     end
 
     it 'limits statement to specified amount of tokens' do
-      lexer = described_class.new("$V1 $V2 $V3 $V4 $V5 $V6")
+      lexer = described_class.new("$V1 $V2 $V3 $V4", max_tokens: 3)
 
-      expect { lexer.tokens(max: 5) }
+      expect { lexer.tokens }
         .to raise_error described_class::SyntaxError
     end