diff --git a/app/assets/javascripts/commons/vue.js b/app/assets/javascripts/commons/vue.js
index cd24a503631352c66c0a3b68456605a0a9833d94..09d2e065e5875d362ae641577ecc43a6c61aa45b 100644
--- a/app/assets/javascripts/commons/vue.js
+++ b/app/assets/javascripts/commons/vue.js
@@ -9,4 +9,4 @@ if (process.env.NODE_ENV !== 'production') {
Vue.use(GlFeatureFlagsPlugin);
Vue.use(Translate);
-Vue.config.ignoredElements = ['gl-emoji'];
+Vue.config.ignoredElements = ['gl-emoji', 'copy-code'];
diff --git a/app/assets/javascripts/diffs/components/app.vue b/app/assets/javascripts/diffs/components/app.vue
index dff95bde26916adf825531e30aef207027f323c1..bfabac3123b7d88519f70e660838acc43b9b5c85 100644
--- a/app/assets/javascripts/diffs/components/app.vue
+++ b/app/assets/javascripts/diffs/components/app.vue
@@ -2,6 +2,7 @@
import { GlLoadingIcon, GlPagination, GlSprintf, GlAlert } from '@gitlab/ui';
import { GlBreakpointInstance as bp } from '@gitlab/ui/dist/utils';
import { mapState, mapGetters, mapActions } from 'vuex';
+import { convertToGraphQLId } from '~/graphql_shared/utils';
import api from '~/api';
import {
keysFor,
@@ -78,6 +79,8 @@ export default {
GlPagination,
GlSprintf,
GlAlert,
+ GenerateTestFileDrawer: () =>
+ import('ee_component/ai/components/generate_test_file_drawer.vue'),
},
mixins: [glFeatureFlagsMixin()],
alerts: {
@@ -226,6 +229,7 @@ export default {
'showWhitespace',
'targetBranchName',
'branchName',
+ 'generateTestFilePath',
]),
...mapGetters('diffs', [
'whichCollapsedTypes',
@@ -297,6 +301,9 @@ export default {
fileReviews() {
return reviewStatuses(this.diffFiles, this.mrReviews);
},
+ resourceId() {
+ return convertToGraphQLId('MergeRequest', this.getNoteableData.id);
+ },
},
watch: {
commit(newCommit, oldCommit) {
@@ -450,6 +457,7 @@ export default {
'navigateToDiffFileIndex',
'setFileByFile',
'disableVirtualScroller',
+ 'setGenerateTestFilePath',
]),
...mapActions('findingsDrawer', ['setDrawer']),
closeDrawer() {
@@ -807,5 +815,11 @@ export default {
</div>
</div>
</div>
+ <generate-test-file-drawer
+ v-if="getNoteableData.id"
+ :resource-id="resourceId"
+ :file-path="generateTestFilePath"
+ @close="() => setGenerateTestFilePath('')"
+ />
</div>
</template>
diff --git a/app/assets/javascripts/diffs/components/diff_file_header.vue b/app/assets/javascripts/diffs/components/diff_file_header.vue
index c3a4897ce780777f5ec6d7f70a56086245d0eaae..e202e02fedd32303754619d6d88d88004baee784 100644
--- a/app/assets/javascripts/diffs/components/diff_file_header.vue
+++ b/app/assets/javascripts/diffs/components/diff_file_header.vue
@@ -52,6 +52,11 @@ export default {
compareButtonLabel: __('Compare submodule commit revisions'),
fileModeTooltip: __('File permissions'),
},
+ inject: {
+ showGenerateTestFileButton: {
+ default: false,
+ },
+ },
props: {
discussionPath: {
type: String,
@@ -227,6 +232,7 @@ export default {
'setCurrentFileHash',
'reviewFile',
'setFileCollapsedByUser',
+ 'setGenerateTestFilePath',
]),
handleToggleFile() {
this.$emit('toggleFile');
@@ -412,6 +418,12 @@ export default {
<gl-icon name="ellipsis_v" class="mr-0" />
<span class="sr-only">{{ $options.i18n.optionsDropdownTitle }}</span>
</template>
+ <gl-dropdown-item
+ v-if="showGenerateTestFileButton"
+ @click="setGenerateTestFilePath(diffFile.new_path)"
+ >
+ {{ __('Generate test with AI') }}
+ </gl-dropdown-item>
<gl-dropdown-item
v-if="diffFile.replaced_view_path"
ref="replacedFileButton"
diff --git a/app/assets/javascripts/diffs/index.js b/app/assets/javascripts/diffs/index.js
index ad7182024da241ec4df93f334a9bebf99b0570da..53c27632c4fadbffba1a82ddbd14f5ff33b46f8e 100644
--- a/app/assets/javascripts/diffs/index.js
+++ b/app/assets/javascripts/diffs/index.js
@@ -28,6 +28,7 @@ export default function initDiffsApp(store = notesStore) {
apolloProvider,
provide: {
newCommentTemplatePath: dataset.newCommentTemplatePath,
+ showGenerateTestFileButton: parseBoolean(dataset.showGenerateTestFileButton),
},
data() {
return {
diff --git a/doc/api/graphql/reference/index.md b/doc/api/graphql/reference/index.md
index 312b02356a5ae2a39ddabd35e255f23b6f7f493b..fec5eabdd9f2f94965ea64ee474ac68d55f997b2 100644
--- a/doc/api/graphql/reference/index.md
+++ b/doc/api/graphql/reference/index.md
@@ -973,6 +973,7 @@ Input type: `AiActionInput`
| <a id="mutationaiactionclientmutationid"></a>`clientMutationId` | [`String`](#string) | A unique identifier for the client performing the mutation. |
| <a id="mutationaiactionexplaincode"></a>`explainCode` | [`AiExplainCodeInput`](#aiexplaincodeinput) | Input for explain_code AI action. |
| <a id="mutationaiactionexplainvulnerability"></a>`explainVulnerability` | [`AiExplainVulnerabilityInput`](#aiexplainvulnerabilityinput) | Input for explain_vulnerability AI action. |
+| <a id="mutationaiactiongeneratetestfile"></a>`generateTestFile` | [`GenerateTestFileInput`](#generatetestfileinput) | Input for generate_test_file AI action. |
| <a id="mutationaiactionmarkupformat"></a>`markupFormat` | [`MarkupFormat`](#markupformat) | Indicates the response format. |
| <a id="mutationaiactionsummarizecomments"></a>`summarizeComments` | [`AiSummarizeCommentsInput`](#aisummarizecommentsinput) | Input for summarize_comments AI action. |
| <a id="mutationaiactiontanukibot"></a>`tanukiBot` | [`AiTanukiBotInput`](#aitanukibotinput) | Input for tanuki_bot AI action. |
@@ -26945,6 +26946,15 @@ Represents an escalation rule.
| <a id="escalationruleinputstatus"></a>`status` | [`EscalationRuleStatus!`](#escalationrulestatus) | Status required to prevent the rule from activating. |
| <a id="escalationruleinputusername"></a>`username` | [`String`](#string) | Username of the user to notify. |
+### `GenerateTestFileInput`
+
+#### Arguments
+
+| Name | Type | Description |
+| ---- | ---- | ----------- |
+| <a id="generatetestfileinputfilepath"></a>`filePath` | [`String!`](#string) | File path to generate test files for. |
+| <a id="generatetestfileinputresourceid"></a>`resourceId` | [`AiModelID!`](#aimodelid) | Global ID of the resource to mutate. |
+
### `JiraUsersMappingInputType`
#### Arguments
diff --git a/ee/app/assets/javascripts/ai/components/generate_test_file_drawer.vue b/ee/app/assets/javascripts/ai/components/generate_test_file_drawer.vue
new file mode 100644
index 0000000000000000000000000000000000000000..9ccc6ba8a3fd3ee43ccaf99f572e208528e071c8
--- /dev/null
+++ b/ee/app/assets/javascripts/ai/components/generate_test_file_drawer.vue
@@ -0,0 +1,143 @@
+<script>
+import { GlDrawer, GlBadge, GlSkeletonLoader, GlAlert } from '@gitlab/ui';
+import { convertToGraphQLId } from '~/graphql_shared/utils';
+import SafeHtml from '~/vue_shared/directives/safe_html';
+import { getContentWrapperHeight } from '~/lib/utils/dom_utils';
+import { DRAWER_Z_INDEX } from '~/lib/utils/constants';
+import aiResponseSubscription from 'ee/graphql_shared/subscriptions/ai_completion_response.subscription.graphql';
+import testFileGeneratorMutation from '../graphql/test_file_generator.mutation.graphql';
+
+export default {
+ apollo: {
+ $subscribe: {
+ testFile: {
+ query: aiResponseSubscription,
+ variables() {
+ return {
+ resourceId: this.resourceId,
+ userId: convertToGraphQLId('User', window.gon.current_user_id), // eslint-disable-line @gitlab/require-i18n-strings
+ };
+ },
+ skip() {
+ return !this.opened;
+ },
+ result({ data }) {
+ const responseBody = data.aiCompletionResponse?.responseBody;
+
+ if (responseBody) {
+ const codeBlockRegex = /<pre(.*)><code>(.*)<\/code><\/pre>/gm;
+ const codeBlock = codeBlockRegex.exec(responseBody.replaceAll('\n', '\\n'));
+
+ if (codeBlock) {
+ this.generatedTest = codeBlock[0].replaceAll('\\n', '\n');
+ this.state = '';
+ } else {
+ this.state = 'unable';
+ }
+ }
+ },
+ },
+ },
+ },
+ directives: { SafeHtml },
+ components: {
+ GlDrawer,
+ GlBadge,
+ GlSkeletonLoader,
+ GlAlert,
+ },
+ props: {
+ resourceId: {
+ type: String,
+ required: true,
+ },
+ filePath: {
+ type: String,
+ required: true,
+ },
+ },
+ data() {
+ return {
+ opened: false,
+ state: '',
+ generatedTest: '',
+ };
+ },
+ computed: {
+ drawerHeightOffset() {
+ return getContentWrapperHeight('.content-wrapper');
+ },
+ },
+ watch: {
+ filePath: {
+ handler(newVal) {
+ this.opened = Boolean(newVal);
+
+ if (this.opened) {
+ this.triggerMutation();
+ } else {
+ this.generatedTest = '';
+ this.state = '';
+ }
+ },
+ immediate: true,
+ },
+ },
+ methods: {
+ triggerMutation() {
+ this.state = 'loading';
+
+ this.$apollo.mutate({
+ mutation: testFileGeneratorMutation,
+ variables: {
+ resourceId: this.resourceId,
+ filePath: this.filePath,
+ },
+ });
+ },
+ },
+ DRAWER_Z_INDEX,
+};
+</script>
+
+<template>
+ <gl-drawer
+ :open="opened"
+ :header-height="drawerHeightOffset"
+ :z-index="$options.DRAWER_Z_INDEX"
+ @close="$emit('close')"
+ >
+ <template #title>
+ <div class="gl-display-flex">
+ <h2 class="gl-my-0 gl-font-size-h2 gl-line-height-24">
+ {{ __('AI Generated Test File') }}
+ </h2>
+ <div>
+ <gl-badge variant="info" size="sm" class="gl-ml-3 gl-mt-2">{{
+ __('Experiment')
+ }}</gl-badge>
+ </div>
+ </div>
+ </template>
+ <div>
+ <div class="markdown-code-block gl-relative">
+ <div
+ v-if="state === 'loading'"
+ class="gl-border-1 gl-border-gray-100 gl-border-solid gl-p-4 gl-rounded-base gl-bg-gray-10"
+ data-testid="generate-test-loading-state"
+ >
+ <gl-skeleton-loader :lines="4" />
+ </div>
+ <gl-alert v-else-if="state === 'unable'" :dismissible="false">
+ {{ __('Unable to generate tests for specified file.') }}
+ </gl-alert>
+ <template v-else>
+ <div class="gl-relative markdown-code-block js-markdown-code">
+ <span v-safe-html="generatedTest" data-testid="generate-test-code"></span>
+ <copy-code />
+ </div>
+ </template>
+ </div>
+ </div>
+ </gl-drawer>
+</template>
diff --git a/ee/app/assets/javascripts/ai/graphql/test_file_generator.mutation.graphql b/ee/app/assets/javascripts/ai/graphql/test_file_generator.mutation.graphql
new file mode 100644
index 0000000000000000000000000000000000000000..7c94ce8bf7f183ba4328d94dea221bca2898d443
--- /dev/null
+++ b/ee/app/assets/javascripts/ai/graphql/test_file_generator.mutation.graphql
@@ -0,0 +1,10 @@
+mutation generateTestFile($resourceId: AiModelID!, $filePath: String!) {
+ aiAction(
+ input: {
+ generateTestFile: { resourceId: $resourceId, filePath: $filePath }
+ markupFormat: HTML
+ }
+ ) {
+ errors
+ }
+}
diff --git a/ee/app/assets/javascripts/diffs/store/actions.js b/ee/app/assets/javascripts/diffs/store/actions.js
index adb6b16a13f8dd3a6cb463820c8731deeb059991..78e0e68bbd69de608a323cfa6984088e2dd68dd8 100644
--- a/ee/app/assets/javascripts/diffs/store/actions.js
+++ b/ee/app/assets/javascripts/diffs/store/actions.js
@@ -78,3 +78,6 @@ export const fetchCodequality = ({ commit, state, dispatch }) => {
}
});
};
+
+export const setGenerateTestFilePath = ({ commit }, path) =>
+ commit(types.SET_GENERATE_TEST_FILE_PATH, path);
diff --git a/ee/app/assets/javascripts/diffs/store/modules/diff_state.js b/ee/app/assets/javascripts/diffs/store/modules/diff_state.js
index 3956c1813b3150cc0481ee08cbddbca1c6429e67..05254bec903cfb0a91e640aca85d0ab4606df402 100644
--- a/ee/app/assets/javascripts/diffs/store/modules/diff_state.js
+++ b/ee/app/assets/javascripts/diffs/store/modules/diff_state.js
@@ -5,4 +5,5 @@ export default () => ({
endpointCodequality: '',
codequalityDiff: {},
+ generateTestFilePath: '',
});
diff --git a/ee/app/assets/javascripts/diffs/store/mutation_types.js b/ee/app/assets/javascripts/diffs/store/mutation_types.js
index 4af2ff7f9de043c927273ce4d02334ff715c0035..91a030c5e003619dca37731d358146b5982f69e4 100644
--- a/ee/app/assets/javascripts/diffs/store/mutation_types.js
+++ b/ee/app/assets/javascripts/diffs/store/mutation_types.js
@@ -1,2 +1,3 @@
export const SET_CODEQUALITY_ENDPOINT = 'SET_CODEQUALITY_ENDPOINT';
export const SET_CODEQUALITY_DATA = 'SET_CODEQUALITY_DATA';
+export const SET_GENERATE_TEST_FILE_PATH = 'SET_GENERATE_TEST_FILE_PATH';
diff --git a/ee/app/assets/javascripts/diffs/store/mutations.js b/ee/app/assets/javascripts/diffs/store/mutations.js
index d18311752d195ba0d34ff04399b12395a6e48ff4..9bcb6a99dd7bce81e81c25683bcebcebd4ea5eab 100644
--- a/ee/app/assets/javascripts/diffs/store/mutations.js
+++ b/ee/app/assets/javascripts/diffs/store/mutations.js
@@ -12,4 +12,8 @@ export default {
[types.SET_CODEQUALITY_DATA](state, codequalityDiffData) {
Object.assign(state, { codequalityDiff: codequalityDiffData });
},
+
+ [types.SET_GENERATE_TEST_FILE_PATH](state, path) {
+ state.generateTestFilePath = path;
+ },
};
diff --git a/ee/app/graphql/types/ai/generate_test_file_input_type.rb b/ee/app/graphql/types/ai/generate_test_file_input_type.rb
new file mode 100644
index 0000000000000000000000000000000000000000..14487baeaa7865eb5348123d4a4880f732329147
--- /dev/null
+++ b/ee/app/graphql/types/ai/generate_test_file_input_type.rb
@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module Types
+ module Ai
+ class GenerateTestFileInputType < BaseMethodInputType
+ graphql_name 'GenerateTestFileInput'
+
+ argument :file_path, GraphQL::Types::String,
+ required: true,
+ validates: { allow_blank: false },
+ description: 'File path to generate test files for.'
+ end
+ end
+end
diff --git a/ee/app/helpers/ee/merge_requests_helper.rb b/ee/app/helpers/ee/merge_requests_helper.rb
index 7fef8e8b98989ace1736167181b75d99ab4b8226..b491381b52340da25d75706cbe173ffa3b2d2cf5 100644
--- a/ee/app/helpers/ee/merge_requests_helper.rb
+++ b/ee/app/helpers/ee/merge_requests_helper.rb
@@ -21,7 +21,8 @@ def render_items_list(items, separator = "and")
override :diffs_tab_pane_data
def diffs_tab_pane_data(project, merge_request, params)
super.merge(
- endpoint_codequality: (codequality_mr_diff_reports_project_merge_request_path(project, merge_request, 'json') if project.licensed_feature_available?(:inline_codequality) && merge_request.has_codequality_mr_diff_report?)
+ endpoint_codequality: (codequality_mr_diff_reports_project_merge_request_path(project, merge_request, 'json') if project.licensed_feature_available?(:inline_codequality) && merge_request.has_codequality_mr_diff_report?),
+ show_generate_test_file_button: ::Llm::GenerateTestFileService.new(current_user, merge_request).valid?.to_s
)
end
end
diff --git a/ee/app/models/gitlab_subscriptions/features.rb b/ee/app/models/gitlab_subscriptions/features.rb
index e5f6b2ac92190be16fd50b3ea6b4bc3561f318ab..5c83b261c64cf2a2910db441bb9cfdcdfb973fac 100644
--- a/ee/app/models/gitlab_subscriptions/features.rb
+++ b/ee/app/models/gitlab_subscriptions/features.rb
@@ -198,6 +198,7 @@ class Features
evaluate_group_level_compliance_pipeline
explain_code
external_audit_events
+ generate_test_file
git_abuse_rate_limit
group_ci_cd_analytics
group_level_compliance_dashboard
diff --git a/ee/app/services/llm/execute_method_service.rb b/ee/app/services/llm/execute_method_service.rb
index 3826a14b1d15185e87050f9bfd469b8c42255c64..26adc44ce07a60f9b19e40278f0c89a32cfc3f6e 100644
--- a/ee/app/services/llm/execute_method_service.rb
+++ b/ee/app/services/llm/execute_method_service.rb
@@ -8,7 +8,8 @@ class ExecuteMethodService < BaseService
explain_vulnerability: ::Llm::ExplainVulnerabilityService,
summarize_comments: Llm::GenerateSummaryService,
explain_code: Llm::ExplainCodeService,
- tanuki_bot: Llm::TanukiBotService
+ tanuki_bot: Llm::TanukiBotService,
+ generate_test_file: Llm::GenerateTestFileService
}.freeze
def initialize(user, resource, method, options = {})
diff --git a/ee/app/services/llm/generate_test_file_service.rb b/ee/app/services/llm/generate_test_file_service.rb
new file mode 100644
index 0000000000000000000000000000000000000000..0c184839a84859bdac158c0a61355fd0025f80b2
--- /dev/null
+++ b/ee/app/services/llm/generate_test_file_service.rb
@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module Llm
+ class GenerateTestFileService < BaseService
+ def valid?
+ super &&
+ Feature.enabled?(:generate_test_file, user) &&
+ resource.resource_parent.root_ancestor.licensed_feature_available?(:generate_test_file)
+ end
+
+ private
+
+ def perform
+ ::Llm::CompletionWorker.perform_async(user.id, resource.id, resource.class.name, :generate_test_file, options)
+
+ success
+ end
+ end
+end
diff --git a/ee/config/feature_flags/development/generate_test_file.yml b/ee/config/feature_flags/development/generate_test_file.yml
new file mode 100644
index 0000000000000000000000000000000000000000..6647a5cef0b5ec52fab5959c5461d96797a8972b
--- /dev/null
+++ b/ee/config/feature_flags/development/generate_test_file.yml
@@ -0,0 +1,8 @@
+---
+name: generate_test_file
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/118365
+rollout_issue_url:
+milestone: '16.0'
+type: development
+group: group::code review
+default_enabled: false
diff --git a/ee/lib/gitlab/llm/open_ai/completions/factory.rb b/ee/lib/gitlab/llm/open_ai/completions/factory.rb
index 85c5261dc9b2a10a58a488ddcacc17735acf6910..f0ead004b4e5ce895d0daca4e3f508ba82e753b3 100644
--- a/ee/lib/gitlab/llm/open_ai/completions/factory.rb
+++ b/ee/lib/gitlab/llm/open_ai/completions/factory.rb
@@ -21,6 +21,10 @@ class Factory
tanuki_bot: {
service_class: ::Gitlab::Llm::OpenAi::Completions::TanukiBot,
prompt_class: ::Gitlab::Llm::OpenAi::Templates::TanukiBot
+ },
+ generate_test_file: {
+ service_class: ::Gitlab::Llm::OpenAi::Completions::GenerateTestFile,
+ prompt_class: ::Gitlab::Llm::OpenAi::Templates::GenerateTestFile
}
}.freeze
diff --git a/ee/lib/gitlab/llm/open_ai/completions/generate_test_file.rb b/ee/lib/gitlab/llm/open_ai/completions/generate_test_file.rb
new file mode 100644
index 0000000000000000000000000000000000000000..74af23782f7d136153d13d933d044d70ad76a7dc
--- /dev/null
+++ b/ee/lib/gitlab/llm/open_ai/completions/generate_test_file.rb
@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module Gitlab
+ module Llm
+ module OpenAi
+ module Completions
+ class GenerateTestFile
+ TOTAL_MODEL_TOKEN_LIMIT = 4000
+ OUTPUT_TOKEN_LIMIT = (TOTAL_MODEL_TOKEN_LIMIT * 0.25).to_i.freeze
+
+ def initialize(ai_prompt_class)
+ @ai_prompt_class = ai_prompt_class
+ end
+
+ def execute(user, merge_request, options)
+ return unless user
+ return unless merge_request
+ return unless merge_request.send_to_ai?
+
+ ai_options = ai_prompt_class.get_options(merge_request, options[:file_path])
+ ai_options[:max_tokens] = OUTPUT_TOKEN_LIMIT
+
+ ai_response = Gitlab::Llm::OpenAi::Client.new(user).chat(content: nil, **ai_options)
+
+ ::Gitlab::Llm::OpenAi::ResponseService.new(user, merge_request, ai_response, options: options).execute(
+ Gitlab::Llm::OpenAi::ResponseModifiers::Chat.new
+ )
+ end
+
+ private
+
+ attr_reader :ai_prompt_class
+ end
+ end
+ end
+ end
+end
diff --git a/ee/lib/gitlab/llm/open_ai/templates/generate_test_file.rb b/ee/lib/gitlab/llm/open_ai/templates/generate_test_file.rb
new file mode 100644
index 0000000000000000000000000000000000000000..717c7d076708e888f183c8caf7b1f7543816f274
--- /dev/null
+++ b/ee/lib/gitlab/llm/open_ai/templates/generate_test_file.rb
@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Gitlab
+ module Llm
+ module OpenAi
+ module Templates
+ class GenerateTestFile
+ def self.get_options(merge_request, path)
+ prompt = <<-TEMPLATE
+ Write unit tests for #{path} to ensure its proper functioning but only if the file contains code
+ """
+ #{Gitlab::Llm::OpenAi::Templates::GenerateTestFile.get_diff_file_content(merge_request, path)}
+ """
+ TEMPLATE
+
+ {
+ content: prompt,
+ temperature: 0.2
+ }
+ end
+
+ def self.get_diff_file_content(merge_request, path)
+ file = merge_request.diffs.diff_files.find { |file| file.paths.include?(path) }
+
+ file&.blob&.data
+ end
+ end
+ end
+ end
+ end
+end
diff --git a/ee/spec/frontend/ai/components/generate_test_file_drawer_spec.js b/ee/spec/frontend/ai/components/generate_test_file_drawer_spec.js
new file mode 100644
index 0000000000000000000000000000000000000000..a221b55172bd9ef47403a65dad4d53f161b93d91
--- /dev/null
+++ b/ee/spec/frontend/ai/components/generate_test_file_drawer_spec.js
@@ -0,0 +1,112 @@
+import Vue from 'vue';
+import VueApollo from 'vue-apollo';
+import { GlAlert } from '@gitlab/ui';
+import waitForPromises from 'helpers/wait_for_promises';
+import createMockApollo from 'helpers/mock_apollo_helper';
+import { mountExtended } from 'helpers/vue_test_utils_helper';
+import GenerateTestFileDrawer from 'ee/ai/components/generate_test_file_drawer.vue';
+import aiResponseSubscription from 'ee/graphql_shared/subscriptions/ai_completion_response.subscription.graphql';
+import testFileGeneratorMutation from 'ee/ai/graphql/test_file_generator.mutation.graphql';
+
+Vue.use(VueApollo);
+Vue.config.ignoredElements = ['copy-code'];
+
+let wrapper;
+let subscriptionHandlerMock;
+let mutationHandlerMock;
+
+function createComponent() {
+ const apolloProvider = createMockApollo([
+ [aiResponseSubscription, subscriptionHandlerMock],
+ [testFileGeneratorMutation, mutationHandlerMock],
+ ]);
+
+ wrapper = mountExtended(GenerateTestFileDrawer, {
+ propsData: {
+ resourceId: 'gid://gitlab/MergeRequest/1',
+ filePath: 'index.js',
+ },
+ apolloProvider,
+ });
+}
+
+describe('Generate test file drawer component', () => {
+ beforeEach(() => {
+ window.gon.current_user_id = 1;
+ mutationHandlerMock = jest
+ .fn()
+ .mockResolvedValue({ data: { aiAction: { errors: [], __typename: 'AiActionPayload' } } });
+ subscriptionHandlerMock = jest.fn().mockResolvedValue({
+ data: {
+ aiCompletionResponse: {
+ responseBody: '<pre><code>This is test code</code></pre>',
+ errors: [],
+ },
+ },
+ });
+ });
+
+ afterEach(() => {
+ mutationHandlerMock.mockRestore();
+ subscriptionHandlerMock.mockRestore();
+ });
+
+ it('calls mutation when mounted', () => {
+ createComponent();
+
+ expect(mutationHandlerMock).toHaveBeenCalledWith({
+ filePath: 'index.js',
+ resourceId: 'gid://gitlab/MergeRequest/1',
+ });
+ });
+
+ it('calls subscription', () => {
+ createComponent();
+
+ expect(subscriptionHandlerMock).toHaveBeenCalledWith({
+ resourceId: 'gid://gitlab/MergeRequest/1',
+ userId: 'gid://gitlab/User/1',
+ });
+ });
+
+ it('shows loading state when subscription is loading', () => {
+ createComponent();
+
+ expect(wrapper.findByTestId('generate-test-loading-state').exists()).toBe(true);
+ });
+
+ it('renders returned test from subscription', async () => {
+ createComponent();
+
+ await waitForPromises();
+
+ expect(wrapper.findByTestId('generate-test-code').text()).toContain('This is test code');
+ });
+
+ it('emits close event when closed', () => {
+ createComponent();
+
+ wrapper.find('.gl-drawer-close-button').vm.$emit('click');
+
+ expect(wrapper.emitted().close).toBeDefined();
+ });
+
+ it('renders alert when test could not be generated', async () => {
+ subscriptionHandlerMock = jest.fn().mockResolvedValue({
+ data: {
+ aiCompletionResponse: {
+ responseBody: 'As the file does not contain any code',
+ errors: [],
+ },
+ },
+ });
+
+ createComponent();
+
+ await waitForPromises();
+
+ expect(wrapper.findComponent(GlAlert).text()).toContain(
+ 'Unable to generate tests for specified file.',
+ );
+ });
+});
diff --git a/ee/spec/frontend/diffs/components/app_spec.js b/ee/spec/frontend/diffs/components/app_spec.js
index 5da8ebaef539955297c6b1616e5e0a083fc3cfe8..c179446aa283b4f7c25f5979bb2a17a63ecb9be5 100644
--- a/ee/spec/frontend/diffs/components/app_spec.js
+++ b/ee/spec/frontend/diffs/components/app_spec.js
@@ -14,6 +14,8 @@ const TEST_ENDPOINT = `${TEST_HOST}/diff/endpoint`;
Vue.use(Vuex);
+Vue.config.ignoredElements = ['copy-code'];
+
describe('diffs/components/app', () => {
let store;
let wrapper;
diff --git a/ee/spec/helpers/merge_requests_helper_spec.rb b/ee/spec/helpers/merge_requests_helper_spec.rb
index 0ceded51364216e7b701a13a1af756bbb84a9966..fbbc88093656b47bcf9cc9108998aef1570d21a9 100644
--- a/ee/spec/helpers/merge_requests_helper_spec.rb
+++ b/ee/spec/helpers/merge_requests_helper_spec.rb
@@ -16,4 +16,17 @@
expect(render_items_list(%w(user user1 user2))).to eq("user, user1 and user2")
end
end
+
+ describe '#diffs_tab_pane_data' do
+ it 'returns data' do
+ project = build_stubbed(:project)
+ merge_request = build_stubbed(:merge_request, project: project)
+
+ allow(helper).to receive(:current_user).and_return(build_stubbed(:user))
+
+ expect(helper.diffs_tab_pane_data(project, merge_request, {})).to include({
+ show_generate_test_file_button: 'false'
+ })
+ end
+ end
end
diff --git a/ee/spec/lib/gitlab/llm/open_ai/completions/generate_test_file_spec.rb b/ee/spec/lib/gitlab/llm/open_ai/completions/generate_test_file_spec.rb
new file mode 100644
index 0000000000000000000000000000000000000000..b11dd26105c778875fd26192448e15fc71ecd353
--- /dev/null
+++ b/ee/spec/lib/gitlab/llm/open_ai/completions/generate_test_file_spec.rb
@@ -0,0 +1,77 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Llm::OpenAi::Completions::GenerateTestFile, feature_category: :code_review_workflow do
+ let_it_be(:user) { create(:user) }
+ let_it_be(:project) { create(:project, :public) }
+ let_it_be(:merge_request) { create(:merge_request, source_project: project) }
+
+ let(:template_class) { ::Gitlab::Llm::OpenAi::Templates::GenerateTestFile }
+ let(:ai_template) { { content: 'something', temperature: 0.2 } }
+ let(:content) { "some ai response text" }
+ let(:ai_response) do
+ {
+ choices: [
+ {
+ message: {
+ content: content
+ }
+ }
+ ]
+ }.to_json
+ end
+
+ subject(:generate_test_file) do
+ described_class.new(template_class).execute(user, merge_request, { file_path: 'index.js' })
+ end
+
+ describe "#execute" do
+ context 'with invalid params' do
+ context 'without user' do
+ let(:user) { nil }
+
+ specify { expect(generate_test_file).to be_nil }
+ end
+
+ context 'without merge request' do
+ let_it_be(:merge_request) { nil }
+
+ specify { expect(generate_test_file).to be_nil }
+ end
+ end
+
+ context 'with valid params' do
+ it 'performs the OpenAI request' do
+ expect_next_instance_of(::Gitlab::Llm::OpenAi::Completions::GenerateTestFile) do |completion_service|
+ expect(completion_service).to receive(:execute).with(user, merge_request, { file_path: 'index.js' })
+ .and_call_original
+ end
+
+ expect(Gitlab::Llm::OpenAi::Templates::GenerateTestFile).to receive(:get_options).and_return(ai_template)
+
+ allow_next_instance_of(Gitlab::Llm::OpenAi::Client) do |instance|
+ params = { content: 'something', max_tokens: 1000, temperature: 0.2 }
+ allow(instance).to receive(:chat).with(params).and_return(ai_response)
+ end
+
+ uuid = 'uuid'
+
+ expect(SecureRandom).to receive(:uuid).and_return(uuid)
+
+ data = {
+ id: uuid,
+ model_name: 'MergeRequest',
+ response_body: content,
+ errors: []
+ }
+
+ expect(GraphqlTriggers).to receive(:ai_completion_response).with(
+ user.to_global_id, merge_request.to_global_id, data
+ )
+
+ generate_test_file
+ end
+ end
+ end
+end
diff --git a/ee/spec/lib/gitlab/llm/open_ai/templates/generate_test_file_spec.rb b/ee/spec/lib/gitlab/llm/open_ai/templates/generate_test_file_spec.rb
new file mode 100644
index 0000000000000000000000000000000000000000..83d1a3c794425f27b075cc1d180fcba23b3208c5
--- /dev/null
+++ b/ee/spec/lib/gitlab/llm/open_ai/templates/generate_test_file_spec.rb
@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Llm::OpenAi::Templates::GenerateTestFile, feature_category: :code_review_workflow do
+ let_it_be(:merge_request) { create(:merge_request) }
+
+ let(:path) { "files/js/commit.coffee" }
+
+ subject { described_class.get_options(merge_request, path) }
+
+ describe '.get_options' do
+ it 'returns correct parameters' do
+ expect(subject[:content]).to include("class Commit")
+ expect(subject[:content]).to include("Write unit tests for #{path} to ensure its proper functioning")
+ expect(subject[:content]).to include("but only if the file contains code")
+ expect(subject[:temperature]).to eq(0.2)
+ end
+ end
+end
diff --git a/ee/spec/requests/api/graphql/mutations/projects/generate_test_file_spec.rb b/ee/spec/requests/api/graphql/mutations/projects/generate_test_file_spec.rb
new file mode 100644
index 0000000000000000000000000000000000000000..3ed698de156659c57a7ae9c60d12a0f36281d16b
--- /dev/null
+++ b/ee/spec/requests/api/graphql/mutations/projects/generate_test_file_spec.rb
@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+
+RSpec.describe 'AiAction for Generate Test File', feature_category: :code_review_workflow do
+ include GraphqlHelpers
+ include Graphql::Subscriptions::Notes::Helper
+
+ let_it_be(:project) { create(:project, :public) }
+ let_it_be(:current_user) { create(:user, developer_projects: [project]) }
+ let_it_be(:merge_request) { create(:merge_request, source_project: project) }
+ let_it_be(:file_path) { "files/js/commit.coffee" }
+
+ let(:mutation) do
+ params = { generate_test_file: { file_path: file_path, resource_id: merge_request.to_gid } }
+
+ graphql_mutation(:ai_action, params) do
+ <<-QL.strip_heredoc
+ errors
+ QL
+ end
+ end
+
+ before do
+ stub_licensed_features(generate_test_file: true)
+ end
+
+ it 'successfully performs an explain code request' do
+ expect(Llm::CompletionWorker).to receive(:perform_async).with(
+ current_user.id, merge_request.id, "MergeRequest", :generate_test_file, {
+ file_path: file_path, markup_format: :raw
+ }
+ )
+
+ post_graphql_mutation(mutation, current_user: current_user)
+
+ expect(graphql_mutation_response(:ai_action)['errors']).to eq([])
+ end
+
+ context 'when empty messages are passed' do
+ let(:file_path) { "" }
+
+ it 'returns nil' do
+ expect(Llm::CompletionWorker).not_to receive(:perform_async)
+
+ post_graphql_mutation(mutation, current_user: current_user)
+
+ expect(fresh_response_data['errors'][0]['message']).to eq("filePath can't be blank")
+ end
+ end
+
+ context 'when openai_experimentation feature flag is disabled' do
+ before do
+ stub_feature_flags(openai_experimentation: false)
+ end
+
+ it 'returns nil' do
+ expect(Llm::CompletionWorker).not_to receive(:perform_async)
+
+ post_graphql_mutation(mutation, current_user: current_user)
+
+ expect(fresh_response_data['errors'][0]['message']).to eq("`openai_experimentation` feature flag is disabled.")
+ end
+ end
+end
diff --git a/ee/spec/services/llm/generate_test_file_service_spec.rb b/ee/spec/services/llm/generate_test_file_service_spec.rb
new file mode 100644
index 0000000000000000000000000000000000000000..586ca57dfdca5c7fa3d81071baad5a40ec051977
--- /dev/null
+++ b/ee/spec/services/llm/generate_test_file_service_spec.rb
@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Llm::GenerateTestFileService, feature_category: :code_review_workflow do
+ let_it_be(:user) { create(:user) }
+ let_it_be(:project) { create(:project, :public) }
+ let_it_be(:merge_request) { create(:merge_request, source_project: project) }
+ let(:options) { {} }
+
+ subject { described_class.new(user, merge_request, options) }
+
+ describe '#execute' do
+ before do
+ stub_licensed_features(generate_test_file: true)
+ allow(Llm::CompletionWorker).to receive(:perform_async)
+ end
+
+ context 'when the user is permitted to view the merge request' do
+ before do
+ project.add_maintainer(user)
+ end
+
+ it 'schedules a job' do
+ expect(subject.execute).to be_success
+
+ expect(Llm::CompletionWorker).to have_received(:perform_async).with(
+ user.id,
+ merge_request.id,
+ 'MergeRequest',
+ :generate_test_file,
+ options
+ )
+ end
+ end
+
+ context 'when the user is not permitted to view the merge request' do
+ it 'returns an error' do
+ project.team.truncate
+
+ expect(subject.execute).to be_error
+
+ expect(Llm::CompletionWorker).not_to have_received(:perform_async)
+ end
+ end
+
+ context 'when feature flag is disabled' do
+ before do
+ stub_feature_flags(generate_test_file: false)
+ end
+
+ it 'returns an error' do
+ expect(subject.execute).to be_error
+
+ expect(Llm::CompletionWorker).not_to have_received(:perform_async)
+ end
+ end
+ end
+end
diff --git a/locale/gitlab.pot b/locale/gitlab.pot
index 4887fed6ca1771888e5615fcfe546b32ddd82c3b..bde74d91365d93eb7201ade233f578a64a85b6d7 100644
--- a/locale/gitlab.pot
+++ b/locale/gitlab.pot
@@ -1833,6 +1833,9 @@ msgstr ""
msgid "ACTION REQUIRED: Something went wrong while obtaining the Let's Encrypt certificate for GitLab Pages domain '%{domain}'"
msgstr ""
+msgid "AI Generated Test File"
+msgstr ""
+
msgid "AI actions"
msgstr ""
@@ -18896,6 +18899,9 @@ msgstr ""
msgid "Generate site and private keys at"
msgstr ""
+msgid "Generate test with AI"
+msgstr ""
+
msgid "Generated with JSON data"
msgstr ""
@@ -47076,6 +47082,9 @@ msgstr ""
msgid "Unable to generate new instance ID"
msgstr ""
+msgid "Unable to generate tests for specified file."
+msgstr ""
+
msgid "Unable to load commits. Try again later."
msgstr ""
diff --git a/spec/frontend/diffs/components/app_spec.js b/spec/frontend/diffs/components/app_spec.js
index e58082a798f748e0946f4d5435a12ce0cf5c2313..aa74721ee3cb01e98ef9d33edf9a6b02e826b1ea 100644
--- a/spec/frontend/diffs/components/app_spec.js
+++ b/spec/frontend/diffs/components/app_spec.js
@@ -31,6 +31,8 @@ const UPDATED_COMMIT_URL = `${TEST_HOST}/COMMIT/NEW`;
Vue.use(Vuex);
+Vue.config.ignoredElements = ['copy-code'];
+
function getCollapsedFilesWarning(wrapper) {
return wrapper.findComponent(CollapsedFilesWarning);
}