mirror of
https://github.com/carlrobertoh/ProxyAI.git
synced 2026-05-12 14:10:29 +00:00
84 lines
3.9 KiB
Java
84 lines
3.9 KiB
Java
package ee.carlrobert.codegpt.codecompletions;
|
|
|
|
import static ee.carlrobert.codegpt.CodeGPTKeys.MULTI_LINE_INLAY;
|
|
import static ee.carlrobert.codegpt.CodeGPTKeys.SINGLE_LINE_INLAY;
|
|
import static ee.carlrobert.codegpt.codecompletions.CodeCompletionService.APPLY_INLAY_ACTION_ID;
|
|
import static ee.carlrobert.codegpt.util.file.FileUtil.getResourceContent;
|
|
import static ee.carlrobert.llm.client.util.JSONUtil.e;
|
|
import static ee.carlrobert.llm.client.util.JSONUtil.jsonMapResponse;
|
|
import static java.util.concurrent.TimeUnit.SECONDS;
|
|
import static org.assertj.core.api.Assertions.assertThat;
|
|
import static org.awaitility.Awaitility.await;
|
|
|
|
import com.intellij.openapi.editor.Editor;
|
|
import com.intellij.openapi.editor.VisualPosition;
|
|
import com.intellij.openapi.util.TextRange;
|
|
import com.intellij.testFramework.PlatformTestUtil;
|
|
import ee.carlrobert.codegpt.settings.configuration.ConfigurationState;
|
|
import ee.carlrobert.llm.client.http.exchange.StreamHttpExchange;
|
|
import java.util.List;
|
|
import testsupport.IntegrationTest;
|
|
|
|
public class CodeCompletionServiceTest extends IntegrationTest {
|
|
|
|
private final VisualPosition cursorPosition = new VisualPosition(3, 0);
|
|
|
|
public void testFetchCodeCompletionLlama() {
|
|
useLlamaService();
|
|
ConfigurationState.getInstance().setCodeCompletionsEnabled(true);
|
|
myFixture.configureByText(
|
|
"CompletionTest.java",
|
|
getResourceContent("/codecompletions/code-completion-file.txt"));
|
|
Editor editor = myFixture.getEditor();
|
|
var expectedCompletion = "TEST_SINGLE_LINE_OUTPUT\nTEST_MULTI_LINE_OUTPUT";
|
|
var prefix = "z".repeat(1015) + "\n[INPUT]\n"; // 512 tokens
|
|
var suffix = "\n[\\INPUT]\n" + "z".repeat(1015); // 512 tokens
|
|
expectLlama((StreamHttpExchange) request -> {
|
|
assertThat(request.getUri().getPath()).isEqualTo("/completion");
|
|
assertThat(request.getMethod()).isEqualTo("POST");
|
|
assertThat(request.getBody())
|
|
.extracting("prompt")
|
|
.isEqualTo(InfillPromptTemplate.LLAMA.buildPrompt(prefix, suffix));
|
|
return List.of(jsonMapResponse(e("content", expectedCompletion), e("stop", true)));
|
|
});
|
|
|
|
editor.getCaretModel().moveToVisualPosition(cursorPosition);
|
|
|
|
await().pollInSameThread().atMost(5, SECONDS)
|
|
.until(() -> {
|
|
PlatformTestUtil.dispatchAllInvocationEventsInIdeEventQueue();
|
|
var singleLineInlayElement = editor.getUserData(SINGLE_LINE_INLAY);
|
|
var multiLineInlayElement = editor.getUserData(MULTI_LINE_INLAY);
|
|
if (singleLineInlayElement != null && multiLineInlayElement != null) {
|
|
var singleLine =
|
|
((InlayInlineElementRenderer) singleLineInlayElement.getRenderer()).getInlayText();
|
|
var multiLine =
|
|
((InlayBlockElementRenderer) multiLineInlayElement.getRenderer()).getInlayText();
|
|
return "TEST_SINGLE_LINE_OUTPUT".equals(singleLine)
|
|
&& "TEST_MULTI_LINE_OUTPUT".equals(multiLine);
|
|
}
|
|
return false;
|
|
});
|
|
}
|
|
|
|
public void testApplyInlayAction() {
|
|
ConfigurationState.getInstance().setAutoFormattingEnabled(false);
|
|
myFixture.configureByText(
|
|
"CompletionTest.java",
|
|
getResourceContent("/codecompletions/code-completion-file.txt"));
|
|
var editor = myFixture.getEditor();
|
|
editor.getCaretModel().moveToVisualPosition(cursorPosition);
|
|
var expectedSingleLineInlay = "FIRST_LINE";
|
|
var expectedMultiLineInlay = "SECOND_LINE\nTHIRD_LINE";
|
|
var expectedInlay = expectedSingleLineInlay + "\n" + expectedMultiLineInlay;
|
|
int cursorOffsetBeforeApply = editor.getCaretModel().getOffset();
|
|
CodeCompletionService.getInstance(getProject())
|
|
.addInlays(editor, cursorOffsetBeforeApply, expectedInlay);
|
|
|
|
myFixture.performEditorAction(APPLY_INLAY_ACTION_ID);
|
|
|
|
var newTextRange = new TextRange(cursorOffsetBeforeApply, editor.getCaretModel().getOffset());
|
|
var appliedInlay = editor.getDocument().getText(newTextRange);
|
|
assertThat(appliedInlay).isEqualTo(expectedInlay);
|
|
}
|
|
}
|