1 files changed, 20 insertions, 7 deletions
diff --git a/spec/frontend/vue_shared/components/source_viewer/utils_spec.js b/spec/frontend/vue_shared/components/source_viewer/utils_spec.js
index 937c3b26c67..0631e7efd54 100644
--- a/spec/frontend/vue_shared/components/source_viewer/utils_spec.js
+++ b/spec/frontend/vue_shared/components/source_viewer/utils_spec.js
@@ -2,12 +2,25 @@ import { wrapLines } from '~/vue_shared/components/source_viewer/utils';
 
 describe('Wrap lines', () => {
   it.each`
-    input                                              | output
-    ${'line 1'}                                        | ${'<span id="LC1" class="line">line 1</span>'}
-    ${'line 1\nline 2'}                                | ${`<span id="LC1" class="line">line 1</span>\n<span id="LC2" class="line">line 2</span>`}
-    ${'<span class="hljs-code">line 1\nline 2</span>'} | ${`<span id="LC1" class="hljs-code">line 1\n<span id="LC2" class="line">line 2</span></span>`}
-    ${'<span class="hljs-code">```bash'}               | ${'<span id="LC1" class="hljs-code">```bash'}
-  `('returns lines wrapped in spans containing line numbers', ({ input, output }) => {
-    expect(wrapLines(input)).toBe(output);
+    content                                            | language             | output
+    ${'line 1'}                                        | ${'javascript'}      | ${'<span id="LC1" lang="javascript" class="line">line 1</span>'}
+    ${'line 1\nline 2'}                                | ${'html'}            | ${`<span id="LC1" lang="html" class="line">line 1</span>\n<span id="LC2" lang="html" class="line">line 2</span>`}
+    ${'<span class="hljs-code">line 1\nline 2</span>'} | ${'html'}            | ${`<span id="LC1" lang="html" class="hljs-code">line 1\n<span id="LC2" lang="html" class="line">line 2</span></span>`}
+    ${'<span class="hljs-code">```bash'}               | ${'bash'}            | ${'<span id="LC1" lang="bash" class="hljs-code">```bash'}
+    ${'<span class="hljs-code">```bash'}               | ${'valid-language1'} | ${'<span id="LC1" lang="valid-language1" class="hljs-code">```bash'}
+    ${'<span class="hljs-code">```bash'}               | ${'valid_language2'} | ${'<span id="LC1" lang="valid_language2" class="hljs-code">```bash'}
+  `('returns lines wrapped in spans containing line numbers', ({ content, language, output }) => {
+    expect(wrapLines(content, language)).toBe(output);
+  });
+
+  it.each`
+    language
+    ${'invalidLanguage>'}
+    ${'"invalidLanguage"'}
+    ${'<invalidLanguage'}
+  `('returns lines safely without XSS language is not valid', ({ language }) => {
+    expect(wrapLines('<span class="hljs-code">```bash', language)).toBe(
+      '<span id="LC1" lang="" class="hljs-code">```bash',
+    );
   });
 });