// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.services.codewhisperer

import com.github.tomakehurst.wiremock.client.WireMock
import com.github.tomakehurst.wiremock.client.WireMock.matching
import com.github.tomakehurst.wiremock.client.WireMock.post
import com.github.tomakehurst.wiremock.client.WireMock.postRequestedFor
import com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo
import com.github.tomakehurst.wiremock.client.WireMock.verify
import com.github.tomakehurst.wiremock.core.WireMockConfiguration
import com.github.tomakehurst.wiremock.junit.WireMockRule
import com.intellij.openapi.application.ApplicationManager
import com.intellij.testFramework.ApplicationRule
import com.intellij.testFramework.DisposableRule
import com.intellij.testFramework.RuleChain
import com.intellij.testFramework.replaceService
import org.junit.After
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import org.mockito.kotlin.mock
import org.mockito.kotlin.whenever
import software.amazon.awssdk.http.SdkHttpClient
import software.amazon.awssdk.regions.Region
import software.amazon.awssdk.services.codewhisperer.CodeWhispererClient
import software.aws.toolkits.core.utils.tryOrNull
import software.aws.toolkits.jetbrains.core.AwsClientManager
import software.aws.toolkits.jetbrains.core.MockClientManager
import software.aws.toolkits.jetbrains.core.MockClientManagerRule
import software.aws.toolkits.jetbrains.core.credentials.MockCredentialManagerRule
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererSettings
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererEndpointCustomizer

/**
 * If this test class failed, highly likely because the endpoint is overwritten by [CodeWhispererEndpointCustomizer]
 */
class CodeWhispererClientTest {
    val applicationRule = ApplicationRule()
    val mockCredentialManager = MockCredentialManagerRule()
    val mockClientManagerRule = MockClientManagerRule()

    @Rule
    @JvmField
    val ruleChain = RuleChain(applicationRule, mockClientManagerRule, mockCredentialManager)

    @Rule
    @JvmField
    val wireMock = WireMockRule(WireMockConfiguration.wireMockConfig().dynamicPort())

    @Rule
    @JvmField
    val disposable = DisposableRule()

    private lateinit var httpClient: SdkHttpClient
    private lateinit var mockExplorerActionManager: CodeWhispererExplorerActionManager
    private lateinit var mockClient: CodeWhispererClient

    @Before
    fun setUp() {
        MockClientManager.useRealImplementations(disposable.disposable)

        // Endpoint will be override by CodeWhispererEndpointCustomizer if there is one, i.e. will call real one instead of localhost
        mockClient = AwsClientManager.getInstance().createUnmanagedClient(
            mockCredentialManager.createCredentialProvider(),
            Region.US_WEST_2,
            "http://127.0.0.1:${wireMock.port()}"
        )
        wireMock.stubFor(
            post("/")
                .willReturn(
                    WireMock.aResponse().withStatus(200)
                )
        )

        mockExplorerActionManager = mock()
        whenever(mockExplorerActionManager.resolveAccessToken()).thenReturn(CodeWhispererTestUtil.testValidAccessToken)
        ApplicationManager.getApplication().replaceService(CodeWhispererExplorerActionManager::class.java, mockExplorerActionManager, disposable.disposable)
    }

    @After
    fun tearDown() {
        tryOrNull { httpClient.close() }
    }

    @Test
    fun `check GetAccessToken request header`() {
        mockClient.getAccessToken {}
        verify(
            postRequestedFor(urlEqualTo("/"))
                .withoutHeader(CodeWhispererEndpointCustomizer.OPTOUT_KEY_NAME)
                .withoutHeader(CodeWhispererEndpointCustomizer.OPTOUT_KEY_NAME)
        )
    }

    @Test
    fun `check ListRecommendation request header`() {
        mockClient.listRecommendations {}
        // default is opt-in (thus opt-out = false)
        verify(
            postRequestedFor(urlEqualTo("/"))
                .withHeader(CodeWhispererEndpointCustomizer.TOKEN_KEY_NAME, matching(CodeWhispererTestUtil.testValidAccessToken))
                .withHeader(CodeWhispererEndpointCustomizer.OPTOUT_KEY_NAME, matching("false"))
        )

        CodeWhispererSettings.getInstance().toggleMetricOptIn(false)
        mockClient.listRecommendations {}
        verify(
            postRequestedFor(urlEqualTo("/"))
                .withHeader(CodeWhispererEndpointCustomizer.TOKEN_KEY_NAME, matching(CodeWhispererTestUtil.testValidAccessToken))
                .withHeader(CodeWhispererEndpointCustomizer.OPTOUT_KEY_NAME, matching("true"))
        )
    }

    @Test
    fun `check createCodeScan request header`() {
        mockClient.createCodeScan {}
        verify(
            postRequestedFor(urlEqualTo("/"))
                .withHeader(CodeWhispererEndpointCustomizer.TOKEN_KEY_NAME, matching(CodeWhispererTestUtil.testValidAccessToken))
                .withoutHeader(CodeWhispererEndpointCustomizer.OPTOUT_KEY_NAME)

        )
    }

    @Test
    fun `check createUploadUrl request header`() {
        mockClient.createCodeScanUploadUrl {}
        verify(
            postRequestedFor(urlEqualTo("/"))
                .withHeader(CodeWhispererEndpointCustomizer.TOKEN_KEY_NAME, matching(CodeWhispererTestUtil.testValidAccessToken))
                .withoutHeader(CodeWhispererEndpointCustomizer.OPTOUT_KEY_NAME)
        )
    }

    @Test
    fun `check listCodeScanFindings request header`() {
        mockClient.listCodeScanFindings {}
        verify(
            postRequestedFor(urlEqualTo("/"))
                .withHeader(CodeWhispererEndpointCustomizer.TOKEN_KEY_NAME, matching(CodeWhispererTestUtil.testValidAccessToken))
                .withoutHeader(CodeWhispererEndpointCustomizer.OPTOUT_KEY_NAME)
        )
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.services.codewhisperer

import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.editor.Document
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.editor.RangeMarker
import com.intellij.openapi.editor.event.DocumentEvent
import com.intellij.openapi.project.Project
import com.intellij.openapi.util.Key
import com.intellij.psi.codeStyle.CodeStyleManager
import com.intellij.testFramework.DisposableRule
import com.intellij.testFramework.fixtures.CodeInsightTestFixture
import com.intellij.testFramework.replaceService
import com.intellij.testFramework.runInEdtAndGet
import com.intellij.testFramework.runInEdtAndWait
import org.assertj.core.api.Assertions.assertThat
import org.junit.After
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import org.mockito.internal.verification.Times
import org.mockito.kotlin.any
import org.mockito.kotlin.argumentCaptor
import org.mockito.kotlin.atLeastOnce
import org.mockito.kotlin.doNothing
import org.mockito.kotlin.doReturn
import org.mockito.kotlin.mock
import org.mockito.kotlin.spy
import org.mockito.kotlin.verify
import org.mockito.kotlin.whenever
import software.aws.toolkits.core.telemetry.MetricEvent
import software.aws.toolkits.core.telemetry.TelemetryBatcher
import software.aws.toolkits.core.telemetry.TelemetryPublisher
import software.aws.toolkits.jetbrains.core.MockClientManagerRule
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonFileName
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponse
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonTestLeftContext
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererLoginType
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJavaScript
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJsx
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPython
import software.aws.toolkits.jetbrains.services.codewhisperer.model.DetailContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.FileContextInfo
import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.RecommendationContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.SupplementalContextInfo
import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererPopupManager.Companion.CODEWHISPERER_USER_ACTION_PERFORMED
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupSettings
import software.aws.toolkits.jetbrains.services.codewhisperer.service.RequestContext
import software.aws.toolkits.jetbrains.services.codewhisperer.service.ResponseContext
import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeCoverageTokens
import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererCodeCoverageTracker
import software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow.CodeWhispererCodeReferenceManager
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.TOTAL_SECONDS_IN_MINUTE
import software.aws.toolkits.jetbrains.services.telemetry.NoOpPublisher
import software.aws.toolkits.jetbrains.services.telemetry.TelemetryService
import software.aws.toolkits.jetbrains.settings.AwsSettings
import software.aws.toolkits.jetbrains.utils.rules.CodeInsightTestFixtureRule
import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule
import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
import software.aws.toolkits.telemetry.CodewhispererCompletionType
import java.util.concurrent.atomic.AtomicInteger

internal abstract class CodeWhispererCodeCoverageTrackerTestBase(myProjectRule: CodeInsightTestFixtureRule) {
    protected class TestCodePercentageTracker(
        timeWindowInSec: Long,
        language: CodeWhispererProgrammingLanguage,
        rangeMarkers: MutableList<RangeMarker> = mutableListOf(),
        codeCoverageTokens: MutableMap<Document, CodeCoverageTokens> = mutableMapOf()
    ) : CodeWhispererCodeCoverageTracker(timeWindowInSec, language, rangeMarkers, codeCoverageTokens, AtomicInteger(0))

    protected class TestTelemetryService(
        publisher: TelemetryPublisher = NoOpPublisher(),
        batcher: TelemetryBatcher
    ) : TelemetryService(publisher, batcher)

    @Rule
    @JvmField
    val projectRule: CodeInsightTestFixtureRule

    @Rule
    @JvmField
    val disposableRule = DisposableRule()

    @Rule
    @JvmField
    val mockClientManagerRule = MockClientManagerRule()

    protected lateinit var project: Project
    protected lateinit var fixture: CodeInsightTestFixture
    protected lateinit var telemetryServiceSpy: TelemetryService
    protected lateinit var batcher: TelemetryBatcher
    protected lateinit var exploreActionManagerMock: CodeWhispererExplorerActionManager

    init {
        this.projectRule = myProjectRule
    }

    open fun setup() {
        this.project = projectRule.project
        this.fixture = projectRule.fixture
        AwsSettings.getInstance().isTelemetryEnabled = true
        batcher = mock()
        telemetryServiceSpy = spy(TestTelemetryService(batcher = batcher))

        exploreActionManagerMock = mock {
            on { checkActiveCodeWhispererConnectionType(any()) } doReturn CodeWhispererLoginType.Sono
        }

        ApplicationManager.getApplication().replaceService(CodeWhispererExplorerActionManager::class.java, exploreActionManagerMock, disposableRule.disposable)
        ApplicationManager.getApplication().replaceService(TelemetryService::class.java, telemetryServiceSpy, disposableRule.disposable)
    }

    @After
    fun tearDown() {
        CodeWhispererCodeCoverageTracker.getInstancesMap().clear()
    }

    protected companion object {
        const val CODE_PERCENTAGE = "codewhisperer_codePercentage"
        const val CWSPR_PERCENTAGE = "codewhispererPercentage"
        const val CWSPR_Language = "codewhispererLanguage"
        const val CWSPR_ACCEPTED_TOKENS = "codewhispererAcceptedTokens"
        const val CWSPR_TOTAL_TOKENS = "codewhispererTotalTokens"
    }
}

internal class CodeWhispererCodeCoverageTrackerTestPython : CodeWhispererCodeCoverageTrackerTestBase(PythonCodeInsightTestFixtureRule()) {
    private lateinit var invocationContext: InvocationContext
    private lateinit var sessionContext: SessionContext

    @Before
    override fun setup() {
        super.setup()
        fixture.configureByText(pythonFileName, pythonTestLeftContext)
        runInEdtAndWait {
            projectRule.fixture.editor.caretModel.primaryCaret.moveToOffset(projectRule.fixture.editor.document.textLength)
        }

        val requestContext = RequestContext(
            project,
            fixture.editor,
            mock(),
            mock(),
            FileContextInfo(mock(), pythonFileName, CodeWhispererPython.INSTANCE),
            SupplementalContextInfo(false, emptyList(), 0, ""),
            null,
            mock()
        )
        val responseContext = ResponseContext("sessionId", CodewhispererCompletionType.Block)
        val recommendationContext = RecommendationContext(
            listOf(DetailContext("requestId", pythonResponse.completions()[0], pythonResponse.completions()[0], false, false)),
            "x, y",
            "x, y",
            mock()
        )
        invocationContext = InvocationContext(requestContext, responseContext, recommendationContext, mock())
        sessionContext = SessionContext()

        // it is needed because referenceManager is listening to CODEWHISPERER_USER_ACTION_PERFORMED topic
        project.replaceService(CodeWhispererCodeReferenceManager::class.java, mock(), disposableRule.disposable)
    }

    @Test
    fun `test getInstance()`() {
        assertThat(CodeWhispererCodeCoverageTracker.getInstancesMap()).hasSize(0)
        val javaInstance = CodeWhispererCodeCoverageTracker.getInstance(CodeWhispererJava.INSTANCE)
        assertThat(CodeWhispererCodeCoverageTracker.getInstancesMap()).hasSize(1)
        assertThat(javaInstance).isNotNull

        val javaInstance2 = CodeWhispererCodeCoverageTracker.getInstance(CodeWhispererJava.INSTANCE)
        assertThat(CodeWhispererCodeCoverageTracker.getInstancesMap()).hasSize(1)
        assertThat(javaInstance == javaInstance2).isTrue

        val pythonInstance = CodeWhispererCodeCoverageTracker.getInstance(CodeWhispererPython.INSTANCE)
        assertThat(CodeWhispererCodeCoverageTracker.getInstancesMap()).hasSize(2)
        val pythonInstance2 = CodeWhispererCodeCoverageTracker.getInstance(CodeWhispererPython.INSTANCE)
        assertThat(pythonInstance == pythonInstance2).isTrue

        CodeWhispererCodeCoverageTracker.getInstance(CodeWhispererJavaScript.INSTANCE)
        assertThat(CodeWhispererCodeCoverageTracker.getInstancesMap()).hasSize(3)
    }

    @Test
    fun `test tracker is listening to cwspr recommendation service invocation`() {
        val pythonTracker = TestCodePercentageTracker(TOTAL_SECONDS_IN_MINUTE, CodeWhispererPython.INSTANCE)
        val jsxTracker = TestCodePercentageTracker(TOTAL_SECONDS_IN_MINUTE, CodeWhispererJsx.INSTANCE)
        CodeWhispererCodeCoverageTracker.getInstancesMap()[CodeWhispererPython.INSTANCE] = pythonTracker
        CodeWhispererCodeCoverageTracker.getInstancesMap()[CodeWhispererJsx.INSTANCE] = jsxTracker
        pythonTracker.activateTrackerIfNotActive()
        assertThat(pythonTracker.serviceInvocationCount).isEqualTo(0)
        assertThat(jsxTracker.serviceInvocationCount).isEqualTo(0)

        val fileContextInfo = mock<FileContextInfo> { on { programmingLanguage } doReturn CodeWhispererPython.INSTANCE }
        ApplicationManager.getApplication().messageBus.syncPublisher(CodeWhispererService.CODEWHISPERER_CODE_COMPLETION_PERFORMED).onSuccess(fileContextInfo)
        assertThat(pythonTracker.serviceInvocationCount).isEqualTo(1)
        assertThat(jsxTracker.serviceInvocationCount).isEqualTo(0)
    }

    @Test
    fun `test tracker is listening to document changes and increment totalTokens - add new code`() {
        val pythonTracker = spy(TestCodePercentageTracker(TOTAL_SECONDS_IN_MINUTE, CodeWhispererPython.INSTANCE))
        CodeWhispererCodeCoverageTracker.getInstancesMap()[CodeWhispererPython.INSTANCE] = pythonTracker
        pythonTracker.activateTrackerIfNotActive()

        fixture.configureByText(pythonFileName, "")
        runInEdtAndWait {
            WriteCommandAction.runWriteCommandAction(project) {
                fixture.editor.appendString(pythonTestLeftContext)
            }
        }
        val captor = argumentCaptor<DocumentEvent>()
        verify(pythonTracker, Times(1)).documentChanged(captor.capture())
        assertThat(captor.firstValue.newFragment.toString()).isEqualTo(pythonTestLeftContext)
        val oldSize = pythonTracker.totalTokensSize
        assertThat(pythonTracker.totalTokensSize).isEqualTo(pythonTestLeftContext.length)

        val anotherCode = "(x, y):"
        runInEdtAndWait {
            WriteCommandAction.runWriteCommandAction(project) {
                fixture.editor.appendString(anotherCode)
            }
        }
        assertThat(pythonTracker.totalTokensSize).isEqualTo(oldSize + anotherCode.length)
    }

    @Test
    fun `test tracker is listening to document changes and increment totalTokens - delete code`() {
        val pythonTracker = TestCodePercentageTracker(
            TOTAL_SECONDS_IN_MINUTE,
            CodeWhispererPython.INSTANCE,
            codeCoverageTokens = mutableMapOf(fixture.editor.document to CodeCoverageTokens(pythonTestLeftContext.length, 0))
        )

        CodeWhispererCodeCoverageTracker.getInstancesMap()[CodeWhispererPython.INSTANCE] = pythonTracker
        pythonTracker.activateTrackerIfNotActive()
        assertThat(pythonTracker.totalTokensSize).isEqualTo(pythonTestLeftContext.length)

        runInEdtAndWait {
            fixture.editor.caretModel.primaryCaret.moveToOffset(fixture.editor.document.textLength)
            WriteCommandAction.runWriteCommandAction(project) {
                fixture.editor.document.deleteString(fixture.editor.caretModel.offset - 3, fixture.editor.caretModel.offset)
            }
        }

        assertThat(pythonTracker.totalTokensSize).isEqualTo(pythonTestLeftContext.length - 3)
    }

    @Test
    fun `test tracker documentChanged - will not increment tokens on blank string of length greater than 1`() {
        val pythonTracker = TestCodePercentageTracker(
            TOTAL_SECONDS_IN_MINUTE,
            CodeWhispererPython.INSTANCE,
            codeCoverageTokens = mutableMapOf(fixture.editor.document to CodeCoverageTokens(pythonTestLeftContext.length, 0))
        )
        CodeWhispererCodeCoverageTracker.getInstancesMap()[CodeWhispererPython.INSTANCE] = pythonTracker
        pythonTracker.activateTrackerIfNotActive()
        assertThat(pythonTracker.totalTokensSize).isEqualTo(pythonTestLeftContext.length)

        runInEdtAndWait {
            WriteCommandAction.runWriteCommandAction(project) {
                fixture.editor.document.insertString(fixture.editor.caretModel.offset, "\t")
            }
        }

        assertThat(pythonTracker.totalTokensSize).isEqualTo(pythonTestLeftContext.length)
    }

    @Test
    fun `test msg CODEWHISPERER_USER_ACTION_PERFORMED will add rangeMarker in the list`() {
        val pythonTracker = spy(TestCodePercentageTracker(TOTAL_SECONDS_IN_MINUTE, language = CodeWhispererPython.INSTANCE))
        pythonTracker.activateTrackerIfNotActive()
        val rangeMarkerMock = runInEdtAndGet {
            spy(fixture.editor.document.createRangeMarker(0, 3)) {
                on { isValid } doReturn true
            }
        }

        ApplicationManager.getApplication().messageBus.syncPublisher(CODEWHISPERER_USER_ACTION_PERFORMED).afterAccept(
            invocationContext,
            mock(),
            rangeMarkerMock
        )

        assertThat(pythonTracker.acceptedRecommendationsCount).isEqualTo(1)
        val argumentCaptor = argumentCaptor<String>()
        verify(rangeMarkerMock, atLeastOnce()).putUserData(any<Key<String>>(), argumentCaptor.capture())
        assertThat(argumentCaptor.firstValue).isEqualTo(pythonTestLeftContext.substring(0, 3))
    }

    @Test
    fun `test 0 totalTokens will return null`() {
        val javaTracker = spy(TestCodePercentageTracker(TOTAL_SECONDS_IN_MINUTE, language = CodeWhispererJava.INSTANCE))
        CodeWhispererCodeCoverageTracker.getInstancesMap()[CodeWhispererJava.INSTANCE] = javaTracker
        assertThat(javaTracker.percentage).isNull()
    }

    @Test
    fun `test flush() will reset tokens and reschedule next telemetry sending`() {
        val pythonTracker = TestCodePercentageTracker(
            TOTAL_SECONDS_IN_MINUTE,
            CodeWhispererPython.INSTANCE,
            codeCoverageTokens = mutableMapOf(mock<Document>() to CodeCoverageTokens("foobar".length, "bar".length))
        )

        pythonTracker.activateTrackerIfNotActive()
        assertThat(pythonTracker.activeRequestCount()).isEqualTo(1)
        assertThat(pythonTracker.acceptedTokensSize).isEqualTo("bar".length)
        assertThat(pythonTracker.totalTokensSize).isEqualTo("foobar".length)

        pythonTracker.forceTrackerFlush()

        assertThat(pythonTracker.activeRequestCount()).isEqualTo(1)
        assertThat(pythonTracker.acceptedTokensSize).isEqualTo(0)
        assertThat(pythonTracker.totalTokensSize).isEqualTo(0)
    }

    @Test
    fun `test when rangeMarker is not vaild, acceptedToken will not be updated`() {
        // when user delete whole recommendation, rangeMarker will be isValid = false
        val rangeMarkerMock: RangeMarker = mock()
        whenever(rangeMarkerMock.isValid).thenReturn(false)
        val pythonTracker = spy(
            TestCodePercentageTracker(
                TOTAL_SECONDS_IN_MINUTE,
                CodeWhispererPython.INSTANCE,
                mutableListOf(rangeMarkerMock),
            )
        ) {
            onGeneric { getAcceptedTokensDelta(any(), any()) } doReturn 100
        }

        pythonTracker.activateTrackerIfNotActive()
        pythonTracker.forceTrackerFlush()

        verify(pythonTracker, Times(0)).getAcceptedTokensDelta(any(), any())
    }

    @Test
    fun `test flush() will call emitTelemetry automatically schedule next call`() {
        val pythonTracker = spy(
            TestCodePercentageTracker(
                TOTAL_SECONDS_IN_MINUTE,
                CodeWhispererPython.INSTANCE,
            )
        )
        doNothing().whenever(pythonTracker).emitCodeWhispererCodeContribution()

        pythonTracker.activateTrackerIfNotActive()
        assertThat(pythonTracker.activeRequestCount()).isEqualTo(1)
        pythonTracker.forceTrackerFlush()

        verify(pythonTracker, Times(1)).emitCodeWhispererCodeContribution()
        assertThat(pythonTracker.activeRequestCount()).isEqualTo(1)
    }

    @Test
    fun `test emitCodeWhispererCodeContribution`() {
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        val rangeMarkerMock1 = mock<RangeMarker> {
            on { isValid } doReturn true
            on { getUserData(any<Key<String>>()) } doReturn "foo"
            on { document } doReturn fixture.editor.document
        }

        val pythonTracker = spy(
            TestCodePercentageTracker(
                TOTAL_SECONDS_IN_MINUTE,
                CodeWhispererPython.INSTANCE,
                mutableListOf(rangeMarkerMock1),
                mutableMapOf(fixture.editor.document to CodeCoverageTokens(totalTokens = 100, acceptedTokens = 0))
            )
        ) {
            onGeneric { extractRangeMarkerString(any()) } doReturn "fou"
            onGeneric { getAcceptedTokensDelta(any(), any()) } doReturn 99
        }

        pythonTracker.emitCodeWhispererCodeContribution()

        val metricCaptor = argumentCaptor<MetricEvent>()
        verify(batcher, Times(1)).enqueue(metricCaptor.capture())
        CodeWhispererTelemetryTest.assertEventsContainsFieldsAndCount(
            metricCaptor.allValues,
            CODE_PERCENTAGE,
            1,
            CWSPR_PERCENTAGE to "99",
            CWSPR_ACCEPTED_TOKENS to "99",
            CWSPR_TOTAL_TOKENS to "100",
            "codewhispererUserGroup" to userGroup.name,
        )
    }

    @Test
    fun `test flush() won't emit telemetry event when users not enabling telemetry`() {
        AwsSettings.getInstance().isTelemetryEnabled = false
        val pythonTracker = spy(TestCodePercentageTracker(TOTAL_SECONDS_IN_MINUTE, CodeWhispererPython.INSTANCE))
        doNothing().whenever(pythonTracker).emitCodeWhispererCodeContribution()

        pythonTracker.activateTrackerIfNotActive()
        pythonTracker.forceTrackerFlush()

        verify(pythonTracker, Times(0)).emitCodeWhispererCodeContribution()
    }

    @Test
    fun `test getAcceptedTokensDelta()`() {
        val tracker = TestCodePercentageTracker(TOTAL_SECONDS_IN_MINUTE, CodeWhispererPython.INSTANCE)
        var originalRecommendation = "foo"
        var modifiedRecommendation = "fou"
        var delta = tracker.getAcceptedTokensDelta(originalRecommendation, modifiedRecommendation)
        assertThat(delta).isEqualTo(2)

        originalRecommendation = "foo"
        modifiedRecommendation = "f11111oo"
        delta = tracker.getAcceptedTokensDelta(originalRecommendation, modifiedRecommendation)
        assertThat(delta).isEqualTo(3)

        originalRecommendation = "foo"
        modifiedRecommendation = "fo"
        delta = tracker.getAcceptedTokensDelta(originalRecommendation, modifiedRecommendation)
        assertThat(delta).isEqualTo(2)

        originalRecommendation = "helloworld"
        modifiedRecommendation = "HelloWorld"
        delta = tracker.getAcceptedTokensDelta(originalRecommendation, modifiedRecommendation)
        assertThat(delta).isEqualTo("helloworld".length - 2)

        originalRecommendation = "helloworld"
        modifiedRecommendation = "World"
        delta = tracker.getAcceptedTokensDelta(originalRecommendation, modifiedRecommendation)
        assertThat(delta).isEqualTo("helloworld".length - "hello".length - 1)

        originalRecommendation = "CodeWhisperer"
        modifiedRecommendation = "CODE"
        delta = tracker.getAcceptedTokensDelta(originalRecommendation, modifiedRecommendation)
        assertThat(delta).isEqualTo(1)

        originalRecommendation = "CodeWhisperer"
        modifiedRecommendation = "codewhispererISBEST"
        delta = tracker.getAcceptedTokensDelta(originalRecommendation, modifiedRecommendation)
        assertThat(delta).isEqualTo("CodeWhisperer".length - 2)

        val pythonCommentAddedByUser = "\"\"\"we don't count this comment as generated by CodeWhisperer\"\"\"\n"
        originalRecommendation = "x, y):\n\treturn x + y"
        modifiedRecommendation = "x, y):\n$pythonCommentAddedByUser\treturn x + y"
        delta = tracker.getAcceptedTokensDelta(originalRecommendation, modifiedRecommendation)
        assertThat(delta).isEqualTo(originalRecommendation.length)
    }

    @Test
    fun `test flush() won't emit telemetry when users are not editing the document (totalTokens == 0)`() {
        val pythonTracker = TestCodePercentageTracker(TOTAL_SECONDS_IN_MINUTE, CodeWhispererPython.INSTANCE)
        pythonTracker.activateTrackerIfNotActive()
        assertThat(pythonTracker.activeRequestCount()).isEqualTo(1)
        pythonTracker.forceTrackerFlush()
        verify(batcher, Times(0)).enqueue(any())
    }

    private fun Editor.appendString(string: String) {
        val currentOffset = caretModel.primaryCaret.offset
        document.insertString(currentOffset, string)
        caretModel.moveToOffset(currentOffset + string.length)
    }
}

internal class CodeWhispererCodeCoverageTrackerTestJava : CodeWhispererCodeCoverageTrackerTestBase(JavaCodeInsightTestFixtureRule()) {
    @Before
    override fun setup() {
        super.setup()
    }

    @Test
    fun `tracker should not update totalTokens if documentChanged events are fired by code reformatting`() {
        val codeNeedToBeReformatted = """
            class Answer {
                private int knapsack(int[] w, int[] v, int c) {
                int[][] dp = new int[w.length + 1][c + 1];
                    for (int i = 0; i < w.length; i++) {for (int j = 0; j <= c; j++) {
                                   if (j < w[i]) {
                                dp[i + 1][j] = dp[i][j];
                        }
                            else {
                      dp[i + 1][j] = Math.max(dp[i][j], dp[i][j - w[i]] + v[i]);
                            }
                        }
                                 }
                    return                  dp[w.length][c];
                }
            }
        """.trimIndent()
        val file = fixture.configureByText("test.java", codeNeedToBeReformatted)
        val tracker = spy(
            TestCodePercentageTracker(
                TOTAL_SECONDS_IN_MINUTE,
                language = CodeWhispererJava.INSTANCE,
                codeCoverageTokens = mutableMapOf(fixture.editor.document to CodeCoverageTokens(totalTokens = codeNeedToBeReformatted.length))
            )
        )
        CodeWhispererCodeCoverageTracker.getInstancesMap()[CodeWhispererJava.INSTANCE] = tracker
        runInEdtAndWait {
            WriteCommandAction.runWriteCommandAction(project) {
                CodeStyleManager.getInstance(project).reformatText(file, 0, fixture.editor.document.textLength)
            }
        }
        // reformat should fire documentChanged events, but tracker should not update token from these events
        verify(tracker, atLeastOnce()).documentChanged(any())
        assertThat(tracker.totalTokensSize).isEqualTo(codeNeedToBeReformatted.length)

        val formatted = """
            class Answer {
                private int knapsack(int[] w, int[] v, int c) {
                    int[][] dp = new int[w.length + 1][c + 1];
                    for (int i = 0; i < w.length; i++) {
                        for (int j = 0; j <= c; j++) {
                            if (j < w[i]) {
                                dp[i + 1][j] = dp[i][j];
                            } else {
                                dp[i + 1][j] = Math.max(dp[i][j], dp[i][j - w[i]] + v[i]);
                            }
                        }
                    }
                    return dp[w.length][c];
                }
            }
        """.trimIndent()
        assertThat(fixture.editor.document.text.trimEnd()).isEqualTo(formatted)
    }
}


// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.services.codewhisperer

import com.intellij.openapi.application.ApplicationManager
import com.intellij.testFramework.replaceService
import com.intellij.ui.dsl.builder.components.DslLabel
import org.assertj.core.api.Assertions.assertThat
import org.junit.Test
import org.mockito.Mockito
import org.mockito.kotlin.doNothing
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanManager
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererConfigurable
import software.aws.toolkits.resources.message
import javax.swing.JCheckBox
import javax.swing.JComponent
import javax.swing.JLabel

class CodeWhispererConfigurableTest : CodeWhispererTestBase() {

    @Test
    fun `test CodeWhisperer configurable`() {
        val codeScanManagerSpy = Mockito.spy(CodeWhispererCodeScanManager.getInstance(projectRule.project))
        doNothing().`when`(codeScanManagerSpy).addCodeScanUI()
        doNothing().`when`(codeScanManagerSpy).removeCodeScanUI()
        projectRule.project.replaceService(CodeWhispererCodeScanManager::class.java, codeScanManagerSpy, disposableRule.disposable)
        val configurable = CodeWhispererConfigurable(projectRule.project)

        // A workaround to initialize disposable in the DslConfigurableBase since somehow the disposable is
        // not configured in the tests if we don't do this
        configurable.reset()

        assertThat(configurable.id).isEqualTo("aws.codewhisperer")
        val panel = configurable.createPanel()
        mockCodeWhispererEnabledStatus(false)

        val checkboxes = panel.components.filterIsInstance<JCheckBox>()

        assertThat(checkboxes.size).isEqualTo(3)
        assertThat(checkboxes.map { it.text }).containsExactlyInAnyOrder(
            message("aws.settings.codewhisperer.include_code_with_reference"),
            message("aws.settings.codewhisperer.configurable.opt_out.title"),
            message("aws.settings.codewhisperer.automatic_import_adder")
        )

        val comments = panel.components.filterIsInstance<DslLabel>()
        assertThat(comments.size).isEqualTo(4)

        mockCodeWhispererEnabledStatus(false)
        ApplicationManager.getApplication().messageBus.syncPublisher(ToolkitConnectionManagerListener.TOPIC)
            .activeConnectionChanged(null)
        checkboxes.forEach {
            assertThat(it.isEnabled).isFalse
        }
        val firstMessage = panel.components.firstOrNull {
            it is JLabel && it.text == message("aws.settings.codewhisperer.warning")
        }
        assertThat(firstMessage).isNotNull
        assertThat((firstMessage as JComponent).isVisible).isTrue

        mockCodeWhispererEnabledStatus(true)
        ApplicationManager.getApplication().messageBus.syncPublisher(ToolkitConnectionManagerListener.TOPIC)
            .activeConnectionChanged(null)
        checkboxes.forEach {
            assertThat(it.isEnabled).isTrue
        }
        assertThat(firstMessage.isVisible).isFalse
    }
}


// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.services.codewhisperer

import com.intellij.help.impl.HelpManagerImpl
import com.intellij.openapi.help.HelpManager
import com.intellij.testFramework.ApplicationRule
import org.assertj.core.api.Assertions.assertThat
import org.junit.Rule
import org.junit.Test
import software.aws.toolkits.jetbrains.core.help.HelpIds
class CodeWhispererConstantsTest {

    @Rule
    @JvmField
    val application = ApplicationRule()

    @Test
    fun `is codewishperer help uri has correct doc link`() {
        assertThat((HelpManager.getInstance() as HelpManagerImpl).getHelpUrl(HelpIds.CODEWHISPERER_LOGIN_YES_NO.id))
            .isEqualTo("https://docs.aws.amazon.com/toolkit-for-jetbrains/latest/userguide/setup-credentials.html")
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.services.codewhisperer

import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.project.Project
import com.intellij.testFramework.fixtures.CodeInsightTestFixture
import com.intellij.testFramework.runInEdtAndGet
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonFileName
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonTestLeftContext
import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPython
import software.aws.toolkits.jetbrains.services.codewhisperer.language.programmingLanguage
import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule

class CodeWhispererEditorUtilTest {
    @Rule
    @JvmField
    var projectRule = PythonCodeInsightTestFixtureRule()

    private lateinit var project: Project
    private lateinit var fixture: CodeInsightTestFixture

    @Before
    fun setup() {
        project = projectRule.project
        fixture = projectRule.fixture
    }

    @Test
    fun `test getFileContextInfo`() {
        val psiFile = fixture.configureByText(pythonFileName, pythonTestLeftContext)
        val fileContext = runInEdtAndGet {
            fixture.editor.caretModel.moveToOffset(fixture.editor.document.textLength)
            CodeWhispererEditorUtil.getFileContextInfo(fixture.editor, psiFile)
        }

        assertThat(fileContext.filename).isEqualTo(pythonFileName)
        assertThat(fileContext.programmingLanguage).isInstanceOf(CodeWhispererPython::class.java)
        assertThat(fileContext.caretContext.leftFileContext).isEqualTo(pythonTestLeftContext)
        assertThat(fileContext.caretContext.rightFileContext).isEqualTo("")
    }

    /**
     * # function to add 2 numbers
     * def addTwoNumbers(a, b)
     *                  ^
     *                  cursor
     */
    @Test
    fun `test extractCaretContext`() {
        val pythonComment = "# function to add 2 numbers\n"
        val pythonTestRightContext = "(a, b)"
        fixture.configureByText(pythonFileName, pythonComment)
        val caretContext = runInEdtAndGet {
            WriteCommandAction.runWriteCommandAction(project) {
                fixture.editor.caretModel.moveToOffset(fixture.editor.document.textLength)
                fixture.editor.document.insertString(fixture.editor.caretModel.offset, pythonTestLeftContext)

                fixture.editor.caretModel.moveToOffset(fixture.editor.document.textLength)
                fixture.editor.document.insertString(fixture.editor.caretModel.offset, pythonTestRightContext)
            }

            CodeWhispererEditorUtil.extractCaretContext(fixture.editor)
        }

        assertThat(caretContext.leftFileContext).isEqualTo("$pythonComment$pythonTestLeftContext")
        assertThat(caretContext.rightFileContext).isEqualTo(pythonTestRightContext)
        assertThat(caretContext.leftContextOnCurrentLine).isEqualTo(pythonTestLeftContext)
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.services.codewhisperer

import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.project.Project
import com.intellij.testFramework.ApplicationRule
import com.intellij.testFramework.DisposableRule
import com.intellij.testFramework.ProjectRule
import com.intellij.testFramework.replaceService
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import org.mockito.kotlin.any
import org.mockito.kotlin.mock
import org.mockito.kotlin.spy
import org.mockito.kotlin.whenever
import software.amazon.awssdk.services.ssooidc.SsoOidcClient
import software.aws.toolkits.core.utils.test.aString
import software.aws.toolkits.jetbrains.core.MockClientManagerRule
import software.aws.toolkits.jetbrains.core.credentials.ManagedBearerSsoConnection
import software.aws.toolkits.jetbrains.core.credentials.MockToolkitAuthManagerRule
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_URL
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererLoginType
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExploreActionState
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererEnabled

class CodeWhispererExplorerActionManagerTest {
    @JvmField
    @Rule
    val applicationRule = ApplicationRule()

    @JvmField
    @Rule
    val projectRule = ProjectRule()

    @JvmField
    @Rule
    val disposableRule = DisposableRule()

    @JvmField
    @Rule
    val authManager = MockToolkitAuthManagerRule()

    @JvmField
    @Rule
    val mockClientManager = MockClientManagerRule()

    private lateinit var mockManager: CodeWhispererExplorerActionManager
    private lateinit var project: Project
    private lateinit var connectionManager: ToolkitConnectionManager

    @Before
    fun setup() {
        mockClientManager.create<SsoOidcClient>()
        project = projectRule.project
        connectionManager = mock()

        project.replaceService(ToolkitConnectionManager::class.java, connectionManager, disposableRule.disposable)
    }

    /**
     * CheckActiveCodeWhispererConnectionType()
     */
    @Test
    fun `when there is no connection, should return logout`() {
        mockManager = spy()
        whenever(connectionManager.activeConnectionForFeature(any())).thenReturn(null)

        val actual = mockManager.checkActiveCodeWhispererConnectionType(project)
        assertThat(actual).isEqualTo(CodeWhispererLoginType.Logout)
    }

    @Test
    fun `when ToS accepted and there is an accountless token, should return accountless`() {
        mockManager = spy()
        mockManager.loadState(
            // set up accountless token
            CodeWhispererExploreActionState().apply {
                this.token = "foo"
            }
        )

        val actual = mockManager.checkActiveCodeWhispererConnectionType(project)
        assertThat(actual).isEqualTo(CodeWhispererLoginType.Accountless)
    }

    @Test
    fun `when ToS accepted, no accountless token and there is an AWS Builder ID connection, should return Sono`() {
        assertLoginType(SONO_URL, CodeWhispererLoginType.Sono)
    }

    @Test
    fun `when ToS accepted, no accountless token and there is an SSO connection, should return SSO`() {
        assertLoginType(aString(), CodeWhispererLoginType.SSO)
    }

    @Test
    fun `test nullifyAccountlessCredentialIfNeeded`() {
        mockManager = CodeWhispererExplorerActionManager()
        mockManager.loadState(CodeWhispererExploreActionState().apply { this.token = "foo" })

        assertThat(mockManager.state.token)
            .isNotNull
            .isEqualTo("foo")

        mockManager.nullifyAccountlessCredentialIfNeeded()

        assertThat(mockManager.state.token)
            .isNull()
    }

    /**
     * isCodeWhispererEnabled
     * - should return false if loginType == Logout
     * - should return true if loginType == Accountless || Sono || SSO
     */
    @Test
    fun `test isCodeWhispererEnabled`() {
        mockManager = mock()
        ApplicationManager.getApplication().replaceService(CodeWhispererExplorerActionManager::class.java, mockManager, disposableRule.disposable)

        whenever(mockManager.checkActiveCodeWhispererConnectionType(project)).thenReturn(CodeWhispererLoginType.Logout)
        assertThat(isCodeWhispererEnabled(project)).isFalse

        whenever(mockManager.checkActiveCodeWhispererConnectionType(project)).thenReturn(CodeWhispererLoginType.Accountless)
        assertThat(isCodeWhispererEnabled(project)).isTrue

        whenever(mockManager.checkActiveCodeWhispererConnectionType(project)).thenReturn(CodeWhispererLoginType.Sono)
        assertThat(isCodeWhispererEnabled(project)).isTrue

        whenever(mockManager.checkActiveCodeWhispererConnectionType(project)).thenReturn(CodeWhispererLoginType.SSO)
        assertThat(isCodeWhispererEnabled(project)).isTrue
    }

    private fun assertLoginType(startUrl: String, expectedType: CodeWhispererLoginType) {
        mockManager = spy()
        val conn: ManagedBearerSsoConnection = mock()
        whenever(connectionManager.activeConnectionForFeature(any())).thenReturn(conn)
        whenever(conn.startUrl).thenReturn(startUrl)
        whenever(conn.getConnectionSettings()).thenReturn(null)

        val actual = mockManager.checkActiveCodeWhispererConnectionType(project)
        assertThat(actual).isEqualTo(expectedType)
    }
}


// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.services.codewhisperer

import com.intellij.openapi.application.runReadAction
import com.intellij.openapi.project.Project
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.PsiFile
import com.intellij.testFramework.fixtures.JavaCodeInsightTestFixture
import com.intellij.testFramework.runInEdtAndGet
import com.intellij.testFramework.runInEdtAndWait
import kotlinx.coroutines.runBlocking
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Ignore
import org.junit.Rule
import org.junit.Test
import org.mockito.kotlin.any
import org.mockito.kotlin.spy
import org.mockito.kotlin.times
import org.mockito.kotlin.verify
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererCsharp
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererKotlin
import software.aws.toolkits.jetbrains.services.codewhisperer.util.DefaultCodeWhispererFileContextProvider
import software.aws.toolkits.jetbrains.services.codewhisperer.util.FileContextProvider
import software.aws.toolkits.jetbrains.utils.rules.HeavyJavaCodeInsightTestFixtureRule
import software.aws.toolkits.jetbrains.utils.rules.addClass
import software.aws.toolkits.jetbrains.utils.rules.addModule
import software.aws.toolkits.jetbrains.utils.rules.addTestClass

class CodeWhispererFileContextProviderTest {
    @JvmField
    @Rule
    val projectRule = HeavyJavaCodeInsightTestFixtureRule()

    lateinit var sut: DefaultCodeWhispererFileContextProvider

    lateinit var fixture: JavaCodeInsightTestFixture
    lateinit var project: Project

    @Before
    fun setup() {
        fixture = projectRule.fixture
        project = projectRule.project

        sut = FileContextProvider.getInstance(project) as DefaultCodeWhispererFileContextProvider
    }

    @Test
    fun `languages not supporting supplemental context will return empty`() {
        val psiFiles = setupFixture(fixture)
        val psi = psiFiles[0]

        runBlocking {
            var context = aFileContextInfo(CodeWhispererCsharp.INSTANCE)

            assertThat(sut.extractSupplementalFileContextForSrc(psi, context)).isEmpty()
            assertThat(sut.extractSupplementalFileContextForTst(psi, context)).isEmpty()

            context = aFileContextInfo(CodeWhispererKotlin.INSTANCE)
            assertThat(sut.extractSupplementalFileContextForSrc(psi, context)).isEmpty()
            assertThat(sut.extractSupplementalFileContextForTst(psi, context)).isEmpty()
        }
    }

    @Test
    fun `extractFileContext should return correct strings`() {
        val src = """
            public class Main {
                public static void main() {
                    System.out.println("Hello world");
                }
            }
        """.trimIndent()
        val psiFile = fixture.configureByText("Main.java", src)

        val fileContext = runInEdtAndGet {
            fixture.editor.caretModel.moveToOffset(47)
            assertThat(fixture.editor.document.text.substring(0, 47)).isEqualTo(
                """
                  public class Main {
                      public static void main
                """.trimIndent()
            )

            assertThat(fixture.editor.document.text.substring(47)).isEqualTo(
                """
                    () {
                            System.out.println("Hello world");
                        }
                    }
                """.trimIndent()
            )

            sut.extractFileContext(fixture.editor, psiFile)
        }

        assertThat(fileContext.filename).isEqualTo("Main.java")
        assertThat(fileContext.programmingLanguage).isEqualTo(CodeWhispererJava.INSTANCE)
        assertThat(fileContext.caretContext.leftFileContext).isEqualTo(
            """
                public class Main {
                    public static void main
            """.trimIndent()
        )
        assertThat(fileContext.caretContext.rightFileContext).isEqualTo(
            """
                () {
                        System.out.println("Hello world");
                    }
                }
            """.trimIndent()
        )
        assertThat(fileContext.caretContext.leftContextOnCurrentLine).isEqualTo("    public static void main")
    }

    @Test
    fun `test extractCodeChunksFromFiles should read files from file producers to get 60 chunks`() {
        val psiFiles = setupFixture(fixture)
        val virtualFiles = psiFiles.mapNotNull { it.virtualFile }
        val javaMainPsiFile = psiFiles.first()

        val fileProducer1: suspend (PsiFile) -> List<VirtualFile> = { psiFile ->
            listOf(virtualFiles[1])
        }

        val fileProducer2: suspend (PsiFile) -> List<VirtualFile> = { psiFile ->
            listOf(virtualFiles[2])
        }

        val result = runBlocking {
            sut.extractCodeChunksFromFiles(javaMainPsiFile, listOf(fileProducer1, fileProducer2))
        }

        assertThat(result[0].content).isEqualTo(
            """public class UtilClass {
            |    public static int util() {};
            |    public static String util2() {};
            """.trimMargin()
        )

        assertThat(result[1].content).isEqualTo(
            """public class UtilClass {
            |    public static int util() {};
            |    public static String util2() {};
            |    private static void helper() {};
            |    public static final int constant1;
            |    public static final int constant2;
            |    public static final int constant3;
            |}
            """.trimMargin()
        )

        assertThat(result[2].content).isEqualTo(
            """public class MyController {
            |    @Get
            |    public Response getRecommendation(Request: req) {}
            """.trimMargin()
        )

        assertThat(result[3].content).isEqualTo(
            """public class MyController {
            |    @Get
            |    public Response getRecommendation(Request: req) {}
            |}
            """.trimMargin()
        )
    }

    /**
     * - src/
     *     - java/
     *          - Main.java
     *          - Util.java
     *          - controllers/
     *              -MyApiController.java
     * - tst/
     *     - java/
     *          - MainTest.java
     *
     */
    // TODO: fix this test, in test env, psiFile.virtualFile == null @psiGist.getFileData(psiFile) { psiFile -> ... }
    @Ignore
    @Test
    fun `extractSupplementalFileContext from src file should extract src`() {
        val psiFiles = setupFixture(fixture)
        sut = spy(sut)

        runReadAction {
            val fileContext = sut.extractFileContext(fixture.editor, psiFiles[0])

            val supplementalContext = runBlocking { sut.extractSupplementalFileContext(psiFiles[0], fileContext) }
            assertThat(supplementalContext?.contents).isNotNull.isNotEmpty
        }

        runBlocking {
            verify(sut).extractSupplementalFileContextForSrc(any(), any())
            verify(sut, times(0)).extractSupplementalFileContextForTst(any(), any())
        }
    }

    /**
     * - src/
     *     - java/
     *          - Main.java
     *          - Util.java
     *          - controllers/
     *              -MyApiController.java
     * - tst/
     *     - java/
     *          - MainTest.java
     *
     */
    // TODO: unignore this case once we configure isUtgSupport = true
    @Ignore
    @Test
    fun `extractSupplementalFileContext from tst file should extract focal file`() {
        val module = fixture.addModule("main")
        fixture.addClass(module, JAVA_MAIN)

        val psiTestClass = fixture.addTestClass(
            module,
            """
            public class MainTest {}
            """
        )

        val tstFile = psiTestClass.containingFile

        sut = spy(sut)

        runReadAction {
            val fileContext = aFileContextInfo(CodeWhispererJava.INSTANCE)
            val supplementalContext = runBlocking {
                sut.extractSupplementalFileContext(tstFile, fileContext)
            }
            assertThat(supplementalContext?.contents)
                .isNotNull
                .isNotEmpty
                .hasSize(1)

            assertThat(supplementalContext?.contents?.get(0)?.content)
                .isNotNull
                .isEqualTo("UTG\n$JAVA_MAIN")
        }

        runBlocking {
            verify(sut, times(0)).extractSupplementalFileContextForSrc(any(), any())
            verify(sut).extractSupplementalFileContextForTst(any(), any())
        }
    }

    private fun setupFixture(fixture: JavaCodeInsightTestFixture): List<PsiFile> {
        val psiFile1 = fixture.addFileToProject("Main.java", JAVA_MAIN)
        val psiFile2 = fixture.addFileToProject("UtilClass.java", JAVA_UTILCLASS)
        val psiFile3 = fixture.addFileToProject("controllers/MyController.java", JAVA_MY_CROLLTER)
        val psiFile4 = fixture.addFileToProject("helpers/Helper1.java", "Class Helper1 {}")
        val psiFile5 = fixture.addFileToProject("helpers/Helper2.java", "Class Helper2 {}")
        val psiFile6 = fixture.addFileToProject("helpers/Helper3.java", "Class Helper3 {}")
        val testPsiFile = fixture.addFileToProject(
            "test/java/MainTest.java",
            """
            public class MainTest {
                @Before
                public void setup() {}
            }
            """.trimIndent()
        )

        runInEdtAndWait {
            fixture.openFileInEditor(psiFile1.virtualFile)
            fixture.editor.caretModel.moveToOffset(fixture.editor.document.textLength)
        }

        return listOf(psiFile1, psiFile2, psiFile3, testPsiFile, psiFile4, psiFile5, psiFile6)
    }

    companion object {
        private val JAVA_MAIN = """public class Main {
            |    public static void main() {
            |        System.out.println("Hello world");
            |    }
            |}
        """.trimMargin()

        private val JAVA_UTILCLASS = """public class UtilClass {
            |    public static int util() {};
            |    public static String util2() {};
            |    private static void helper() {};
            |    public static final int constant1;
            |    public static final int constant2;
            |    public static final int constant3;
            |}
        """.trimMargin()

        private val JAVA_MY_CROLLTER = """public class MyController {
            |    @Get
            |    public Response getRecommendation(Request: req) {}
            |}
        """.trimMargin()
    }
}


// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.services.codewhisperer

import com.intellij.openapi.application.runReadAction
import com.intellij.openapi.editor.EditorFactory
import com.intellij.openapi.project.Project
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.PsiFile
import com.intellij.testFramework.fixtures.CodeInsightTestFixture
import com.intellij.testFramework.runInEdtAndWait
import kotlinx.coroutines.runBlocking
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import software.aws.toolkits.core.utils.test.aString
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererFileCrawler
import software.aws.toolkits.jetbrains.services.codewhisperer.util.FileCrawler
import software.aws.toolkits.jetbrains.services.codewhisperer.util.JavaCodeWhispererFileCrawler
import software.aws.toolkits.jetbrains.utils.rules.CodeInsightTestFixtureRule
import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule

open class CodeWhispererFileCrawlerTest(projectRule: CodeInsightTestFixtureRule) {
    @JvmField
    @Rule
    val projectRule: CodeInsightTestFixtureRule = projectRule

    lateinit var fixture: CodeInsightTestFixture
    lateinit var project: Project

    open fun setup() {
        fixture = projectRule.fixture
        project = projectRule.project
    }
}

class JavaCodeWhispererFileCrawlerTest : CodeWhispererFileCrawlerTest(JavaCodeInsightTestFixtureRule()) {
    lateinit var sut: CodeWhispererFileCrawler

    @Before
    override fun setup() {
        super.setup()
        sut = JavaCodeWhispererFileCrawler
    }

    @Test
    fun `findFilesUnderProjectRoot`() {
        val mainClass = fixture.addFileToProject(
            "Main.java",
            """
            package com.cw.file_crawler_test;

            import java.util.Map;
            import java.util.regex.Pattern;

            import com.cw.file_crawler_test.utils.AnotherClass;
            import com.cw.file_crawler_test.service.controllers.MyController;

            public class Main {
            };
            """.trimIndent()
        )

        val controllerClass = fixture.addFileToProject(
            "service/controllers/MyController.java",
            """
                package com.cw.file_crawler_test.service.controllers;

                public class MyController {}
            """.trimIndent()
        )

        val anotherClass = fixture.addFileToProject(
            "/utils/AnotherClass.java",
            """
                package com.cw.file_crawler_test.utils;

                public class AnotherClass {}
            """.trimIndent()
        )

        val notImportedClass = fixture.addFileToProject(
            "/utils/NotImported.java",
            """
                package com.cw.file_crawler_test.utils;

                public class NotImported {}
            """.trimIndent()
        )

        val notImportedClass2 = fixture.addFileToProject(
            "/utils/NotImported2.java",
            """
                package com.cw.file_crawler_test.utils;

                public class NotImported2 {}
            """.trimIndent()
        )

        fun assertCrawlerFindCorrectFiles(sut: CodeWhispererFileCrawler) {
            runReadAction {
                val actual = sut.listFilesUnderProjectRoot(project)
                val expected = listOf<PsiFile>(mainClass, controllerClass, anotherClass, notImportedClass, notImportedClass2)
                    .map { it.virtualFile }
                    .toSet()

                assertThat(actual).hasSize(expected.size)
                actual.forEach {
                    assertThat(expected.contains(it)).isTrue
                }
            }
        }

        assertCrawlerFindCorrectFiles(JavaCodeWhispererFileCrawler)
    }

    @Test
    fun `listFilesWithinSamePackage`() {
        val targetFile = fixture.addFileToProject("/utils/AnotherClass.java", "")
        val file2Package1 = fixture.addFileToProject("/utils/NotImported.java", "")
        val file3Package1 = fixture.addFileToProject("/utils/NotImported2.java", "")
        fixture.addFileToProject("Main.java", "")
        fixture.addFileToProject("service/controllers/MyController.java", "")

        runReadAction {
            val actual = sut.listFilesWithinSamePackage(targetFile)
            val expected = listOf<PsiFile>(file2Package1, file3Package1)
                .map { it.virtualFile }
                .toSet()

            assertThat(actual).hasSize(expected.size)
            actual.forEach {
                assertThat(expected.contains(it)).isTrue
            }
        }
    }

    @Test
    fun `findFilesImported`() {
        val mainClass = fixture.addFileToProject(
            "Main.java",
            """
            package com.cw.file_crawler_test;

            import java.util.Map;
            import java.util.regex.Pattern;

            import com.cw.file_crawler_test.utils.AnotherClass;
            import com.cw.file_crawler_test.service.controllers.MyController;

            public class Main {
            };
            """.trimIndent()
        )

        val controllerClass = fixture.addFileToProject(
            "service/controllers/MyController.java",
            """
                package com.cw.file_crawler_test.service.controllers;

                public class MyController {}
            """.trimIndent()
        )

        val anotherClass = fixture.addFileToProject(
            "/utils/AnotherClass.java",
            """
                package com.cw.file_crawler_test.utils;

                public class AnotherClass {}
            """.trimIndent()
        )

        fun assertCrawlerFindCorrectFiles(sut: CodeWhispererFileCrawler) {
            runReadAction {
                val expected = setOf<VirtualFile>(controllerClass.virtualFile, anotherClass.virtualFile)
                val actualFiles = runBlocking { sut.listFilesImported(mainClass) }

                assertThat(actualFiles).hasSize(2)
                actualFiles.forEach {
                    assertThat(expected).contains(it)
                }
            }
        }

        assertCrawlerFindCorrectFiles(JavaCodeWhispererFileCrawler)
        // can't make it work right since the temp file created is not in the real file system
        // Naive crawler will actually read the file system thun unable to find files
        //        assertCrawlerFindCorrectFiles(NaiveJavaCodeWhispererFileCrawler(project))
    }

    @Test
    fun `findFocalFileForTest by name`() {
        val mainPsi = fixture.addFileToProject("Main.java", aString())
        fixture.addFileToProject("Class1.java", aString())
        fixture.addFileToProject("Class2.java", aString())
        fixture.addFileToProject("Class3.java", aString())
        val tstPsi = fixture.addFileToProject("/tst/java/MainTest.java", aString())

        fun assertCrawlerFindCorrectFiles(sut: FileCrawler) {
            runInEdtAndWait {
                fixture.openFileInEditor(tstPsi.virtualFile)

                val actual = sut.findFocalFileForTest(tstPsi)

                assertThat(actual).isNotNull.isEqualTo(mainPsi.virtualFile)
            }
        }

        assertCrawlerFindCorrectFiles(JavaCodeWhispererFileCrawler)
    }

    @Test
    fun `findFocalFileForTest by content`() {
        val mainPsi = fixture.addFileToProject(
            "Main.java",
            """
            public class Main {
                public static void main () {
                    runApp()
                }

                public void runApp() {
                    // TODO
                }
            }
            """.trimIndent()
        )
        val file1 = fixture.addFileToProject("Class1.java", "trivial string 1")
        val file2 = fixture.addFileToProject("Class2.java", "trivial string 2")
        val file3 = fixture.addFileToProject("Class3.java", "trivial string 3")
        val tstPsi = fixture.addFileToProject(
            "/tst/java/NotMatchingFileNameTest.java",
            """
            public class MainTest {
                public void testRunApp() {
                    sut.runApp()
                }
            }
            """.trimIndent()
        )

        runInEdtAndWait {
            fixture.openFileInEditor(mainPsi.virtualFile)
            fixture.openFileInEditor(file1.virtualFile)
            fixture.openFileInEditor(file2.virtualFile)
            fixture.openFileInEditor(file3.virtualFile)
            fixture.openFileInEditor(tstPsi.virtualFile)
        }

        fun assertCrawlerFindCorrectFiles(sut: FileCrawler) {
            runInEdtAndWait {
                val openedFiles = EditorFactory.getInstance().allEditors.size

                val actual = sut.findFocalFileForTest(tstPsi)

                assertThat(openedFiles).isEqualTo(5)
                assertThat(actual).isNotNull.isEqualTo(mainPsi.virtualFile)
            }
        }

        assertCrawlerFindCorrectFiles(JavaCodeWhispererFileCrawler)
    }

    @Test
    fun `test util countSubstringMatches`() {
        val elementsToCheck = listOf("apple", "pineapple", "banana", "chocolate", "fries", "laptop", "amazon", "codewhisperer", "aws")
        val targetElements = listOf(
            "an apple a day, keep doctors away",
            "codewhisperer is the best AI code generator",
            "chocolateCake",
            "green apple is sour",
            "pineapple juice",
            "chocolate cake is good"
        )

        val actual = CodeWhispererFileCrawler.countSubstringMatches(targetElements, elementsToCheck)
        assertThat(actual).isEqualTo(4)
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.services.codewhisperer

import com.intellij.openapi.extensions.ExtensionPointName
import com.intellij.openapi.fileTypes.FileType
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.PsiFile
import com.intellij.testFramework.ApplicationRule
import com.intellij.testFramework.DisposableRule
import org.assertj.core.api.Assertions.assertThat
import org.junit.Rule
import org.junit.Test
import org.mockito.kotlin.doReturn
import org.mockito.kotlin.mock
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererLanguageManager
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererC
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererCpp
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererCsharp
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererGo
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJavaScript
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJsx
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererKotlin
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPhp
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPlainText
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPython
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererRuby
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererRust
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererScala
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererShell
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererSql
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererTsx
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererTypeScript
import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
import software.aws.toolkits.telemetry.CodewhispererLanguage
import kotlin.reflect.full.createInstance
import kotlin.reflect.full.primaryConstructor
import kotlin.reflect.jvm.isAccessible

class CodeWhispererLanguageManagerTest {
    @Rule
    @JvmField
    val applicationRule = ApplicationRule()

    @Rule
    @JvmField
    val projectRule = PythonCodeInsightTestFixtureRule()

    @Rule
    @JvmField
    val disposableRule = DisposableRule()

    val manager = CodeWhispererLanguageManager()

    @Test
    fun `test CodeWhispererProgrammingLanguage should be singleton`() {
        val fileTypeMock = mock<FileType> {
            on { name } doReturn "java"
        }
        val vFileMock = mock<VirtualFile> {
            on { fileType } doReturn fileTypeMock
        }

        val lang1 = manager.getLanguage(vFileMock)
        val lang2 = manager.getLanguage(vFileMock)

        assertThat(lang1).isSameAs(lang2)
    }

    @Test
    fun `test getProgrammingLanguage(virtualFile)`() {
        testGetProgrammingLanguageUtil(listOf("java", "Java", "JAVA"), listOf("java"), CodeWhispererJava::class.java)
        testGetProgrammingLanguageUtil(listOf("python", "Python"), listOf("py"), CodeWhispererPython::class.java)
        testGetProgrammingLanguageUtil(listOf("javascript", "JavaScript"), listOf("js"), CodeWhispererJavaScript::class.java)
        testGetProgrammingLanguageUtil(listOf("jsx harmony"), listOf("jsx"), CodeWhispererJsx::class.java)
        testGetProgrammingLanguageUtil(listOf("typescript jsx"), listOf("tsx"), CodeWhispererTsx::class.java)
        testGetProgrammingLanguageUtil(listOf("typescript", "TypeScript"), listOf("ts"), CodeWhispererTypeScript::class.java)
        testGetProgrammingLanguageUtil(listOf("c#", "C#"), listOf("cs"), CodeWhispererCsharp::class.java)
        testGetProgrammingLanguageUtil(listOf("go", "Go"), listOf("go"), CodeWhispererGo::class.java)
        testGetProgrammingLanguageUtil(listOf("kotlin", "Kotlin"), listOf("kt"), CodeWhispererKotlin::class.java)
        testGetProgrammingLanguageUtil(listOf("php", "Php"), listOf("php"), CodeWhispererPhp::class.java)
        testGetProgrammingLanguageUtil(listOf("ruby", "Ruby"), listOf("rb"), CodeWhispererRuby::class.java)
        testGetProgrammingLanguageUtil(listOf("scala", "Scala"), listOf("scala"), CodeWhispererScala::class.java)
        testGetProgrammingLanguageUtil(listOf("sql", "Sql"), listOf("sql"), CodeWhispererSql::class.java)
        testGetProgrammingLanguageUtil(listOf("plain_text"), listOf("txt"), CodeWhispererPlainText::class.java)
        testGetProgrammingLanguageUtil(listOf("c++"), listOf("cpp", "c++", "cc"), CodeWhispererCpp::class.java)
        testGetProgrammingLanguageUtil(listOf("c"), listOf("c", "h"), CodeWhispererC::class.java)
        testGetProgrammingLanguageUtil(listOf("Shell Script"), listOf("sh"), CodeWhispererShell::class.java)
        testGetProgrammingLanguageUtil(listOf("Rust"), listOf("rs"), CodeWhispererRust::class.java)
    }

    @Test
    fun `psiFile passed to getProgrammingLanguage(psiFile) returns null`() {
        // psiFile.virtualFile potentially will return null if virtualFile only exist in the memory instead of the disk
        val psiFileMock = mock<PsiFile> {
            on { virtualFile } doReturn null
            on { name } doReturn "my_python_script_1.py"
        }
        assertThat(manager.getLanguage(psiFileMock)).isInstanceOf(CodeWhispererPython::class.java)
    }

    private fun <T : CodeWhispererProgrammingLanguage> testGetProgrammingLanguageUtil(
        fileTypeNames: List<String>,
        fileExtensions: List<String?>?,
        expectedLanguage: Class<T>
    ) {
        fileExtensions?.forEach { fileExtension ->
            fileTypeNames.forEach { fileTypeName ->
                val fileTypeMock = mock<FileType> {
                    on { name } doReturn fileTypeName
                }
                val vFileMock = mock<VirtualFile> {
                    on { fileType } doReturn fileTypeMock
                    on { extension } doReturn fileExtension
                }
                assertThat(manager.getLanguage(vFileMock)).isInstanceOf(expectedLanguage)
            }
        }
    }
}

class CodeWhispererProgrammingLanguageTest {
    class TestLanguage : CodeWhispererProgrammingLanguage() {
        override val languageId: String = "test-language"
        override fun toTelemetryType(): CodewhispererLanguage = CodewhispererLanguage.Unknown
    }

    @Test
    fun `test language isSupport`() {
        EP_NAME.extensionList.forEach { language ->
            val telemetryType = language.toTelemetryType()
            val shouldSupportAutoCompletion = true

            val shouldSupportSecurityScan = when (telemetryType) {
                CodewhispererLanguage.Java -> true
                CodewhispererLanguage.Python -> true
                else -> false
            }

            assertThat(language.isCodeCompletionSupported()).isEqualTo(shouldSupportAutoCompletion)
            assertThat(language.isCodeScanSupported()).isEqualTo(shouldSupportSecurityScan)
        }
    }

    @Test
    fun `test CodeWhispererProgrammingLanguage isEqual will compare its languageId`() {
        val instance1: CodeWhispererJava = CodeWhispererJava.INSTANCE
        val instance2: CodeWhispererJava
        CodeWhispererJava::class.apply {
            val constructor = primaryConstructor
            constructor?.isAccessible = true
            instance2 = this.createInstance()
        }

        assertThat(instance1).isNotSameAs(instance2)
        assertThat(instance1).isEqualTo(instance2)
    }

    @Test
    fun `test any class extending CodeWhispererProgrammingLanguage isEqual will compare its languageId`() {
        val instance1: TestLanguage
        val instance2: TestLanguage
        TestLanguage::class.apply {
            val constructor = primaryConstructor
            constructor?.isAccessible = true
            instance1 = this.createInstance()
            instance2 = this.createInstance()
        }

        assertThat(instance1).isNotSameAs(instance2)
        assertThat(instance1).isEqualTo(instance2)
    }

    @Test
    fun `test hashCode`() {
        val set = mutableSetOf<CodeWhispererProgrammingLanguage>()
        val instance1 = CodeWhispererJava.INSTANCE
        val instance2: CodeWhispererProgrammingLanguage
        CodeWhispererJava::class.apply {
            val constructor = primaryConstructor
            constructor?.isAccessible = true
            instance2 = this.createInstance()
        }

        set.add(instance1)
        val flag = set.contains(instance2)
        assertThat(flag).isTrue
    }

    companion object {
        val EP_NAME = ExtensionPointName<CodeWhispererProgrammingLanguage>("aws.toolkit.codewhisperer.programmingLanguage")
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.services.codewhisperer

import com.fasterxml.jackson.databind.DeserializationFeature
import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
import com.fasterxml.jackson.module.kotlin.readValue
import org.assertj.core.api.Assertions.assertThat
import org.junit.Test
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererLicenseInfoManager
import software.aws.toolkits.jetbrains.utils.runUnderProgressIfNeeded
import software.aws.toolkits.resources.message

class CodeWhispererLicenseInfoManagerTest : CodeWhispererTestBase() {
    private val mapper = jacksonObjectMapper().disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
    private val licenseLinks by lazy {
        runUnderProgressIfNeeded(null, message("codewhisperer.loading_licenses"), cancelable = false) {
            this.javaClass.getResourceAsStream("/codewhisperer/licenses.json")?.use { resourceStream ->
                mapper.readValue<Map<String, String>>(resourceStream)
            } ?: throw RuntimeException("CodeWhisperer license info not found")
        }
    }

    @Test
    fun `test service return correct license links for licenses in the file licenses json`() {
        val manager = CodeWhispererLicenseInfoManager.getInstance()
        licenseLinks.forEach { (k, v) ->
            assertThat(manager.getLicenseLink(k)).isEqualTo(v)
        }
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.services.codewhisperer

import com.intellij.openapi.actionSystem.DataContext
import com.intellij.openapi.actionSystem.IdeActions
import com.intellij.openapi.editor.actionSystem.EditorActionManager
import org.assertj.core.api.Assertions.assertThat
import org.junit.Test
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponse
import javax.swing.JButton

class CodeWhispererNavigationTest : CodeWhispererTestBase() {

    @Test
    fun `test navigating to previous recommendation should decrement selected index and update label texts by clicking button`() {
        testNavigation(true)
    }

    @Test
    fun `test navigating to next recommendation should increment selected index and update label texts by clicking button`() {
        testNavigation(false)
    }

    @Test
    fun `test navigating to previous recommendation should decrement selected index and update label texts key shortcuts`() {
        testNavigation(true, useKeyboard = true)
    }

    @Test
    fun `test navigating to next recommendation should increment selected index and update label texts by key shortcuts`() {
        testNavigation(false, useKeyboard = true)
    }

    private fun testNavigation(isReverse: Boolean, useKeyboard: Boolean = false) {
        withCodeWhispererServiceInvokedAndWait {
            val indexChange = if (isReverse) -1 else 1

            assertThat(popupManagerSpy.sessionContext.selectedIndex).isEqualTo(0)

            val expectedCount = pythonResponse.completions().size
            var expectedSelectedIndex: Int
            val navigationButton: JButton
            val oppositeButton: JButton
            if (isReverse) {
                navigationButton = popupManagerSpy.popupComponents.prevButton
                oppositeButton = popupManagerSpy.popupComponents.nextButton
                expectedSelectedIndex = expectedCount - 1
            } else {
                navigationButton = popupManagerSpy.popupComponents.nextButton
                oppositeButton = popupManagerSpy.popupComponents.prevButton
                expectedSelectedIndex = 0
            }
            if (isReverse) {
                repeat(expectedCount - 1) {
                    navigateHelper(false, useKeyboard)
                }
            }

            assertThat(popupManagerSpy.sessionContext.selectedIndex).isEqualTo(expectedSelectedIndex)
            assertThat(oppositeButton.isEnabled).isEqualTo(false)

            repeat(expectedCount - 1) {
                assertThat(navigationButton.isEnabled).isEqualTo(true)
                navigateHelper(isReverse, useKeyboard)
                assertThat(oppositeButton.isEnabled).isEqualTo(true)
                expectedSelectedIndex = (expectedSelectedIndex + indexChange) % expectedCount
                assertThat(popupManagerSpy.sessionContext.selectedIndex).isEqualTo(expectedSelectedIndex)
                checkRecommendationInfoLabelText(expectedSelectedIndex + 1, expectedCount)
            }
            assertThat(navigationButton.isEnabled).isEqualTo(false)
        }
    }

    private fun navigateHelper(isReverse: Boolean, useKeyboard: Boolean) {
        if (useKeyboard) {
            val actionHandler = EditorActionManager.getInstance()
            if (isReverse) {
                val leftArrowHandler = actionHandler.getActionHandler(IdeActions.ACTION_EDITOR_MOVE_CARET_LEFT)
                leftArrowHandler.execute(projectRule.fixture.editor, null, DataContext.EMPTY_CONTEXT)
            } else {
                val rightArrowHandler = actionHandler.getActionHandler(IdeActions.ACTION_EDITOR_MOVE_CARET_RIGHT)
                rightArrowHandler.execute(projectRule.fixture.editor, null, DataContext.EMPTY_CONTEXT)
            }
        } else {
            if (isReverse) {
                popupManagerSpy.popupComponents.prevButton.doClick()
            } else {
                popupManagerSpy.popupComponents.nextButton.doClick()
            }
        }
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.services.codewhisperer

import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.project.Project
import com.intellij.testFramework.DisposableRule
import com.intellij.testFramework.fixtures.CodeInsightTestFixture
import com.intellij.testFramework.runInEdtAndGet
import com.intellij.testFramework.runInEdtAndWait
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import software.amazon.awssdk.services.codewhispererruntime.model.Reference
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererRecommendationManager
import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule

class CodeWhispererRecommendationManagerTest {
    @Rule
    @JvmField
    var projectRule = PythonCodeInsightTestFixtureRule()

    @Rule
    @JvmField
    val disposableRule = DisposableRule()

    private val documentContentContent = "012345678"
    private lateinit var fixture: CodeInsightTestFixture
    private lateinit var project: Project
    private val originalReference = Reference.builder()
        .licenseName("test_license")
        .repository("test_repo")
        .build()

    @Before
    fun setup() {
        fixture = projectRule.fixture
        project = projectRule.project

        fixture.configureByText("test.py", documentContentContent)
        runInEdtAndWait {
            fixture.editor.caretModel.moveToOffset(documentContentContent.length)
        }
    }

    @Test
    fun `test reformatReference() should generate a new reference with span based on rangeMarker and no surfix newline char`() {
        // invocationOffset and markerStartOffset is of our choice as long as invocationOffset <= markerStartOffset
        val recommendationManager = CodeWhispererRecommendationManager()
        testReformatReferenceUtil(recommendationManager, documentContentSurfix = "", invocationOffset = 2, markerStartOffset = 5)
        testReformatReferenceUtil(recommendationManager, documentContentSurfix = "\n", invocationOffset = 2, markerStartOffset = 5)
        testReformatReferenceUtil(recommendationManager, documentContentSurfix = "\n\n", invocationOffset = 1, markerStartOffset = 4)
    }

    private fun testReformatReferenceUtil(
        recommendationManager: CodeWhispererRecommendationManager,
        documentContentSurfix: String,
        invocationOffset: Int,
        markerStartOffset: Int
    ) {
        // insert newline characters
        WriteCommandAction.runWriteCommandAction(project) {
            fixture.editor.document.insertString(fixture.editor.caretModel.offset, documentContentSurfix)
        }

        val rangeMarker =
            runInEdtAndGet { fixture.editor.document.createRangeMarker(markerStartOffset, documentContentContent.length + documentContentSurfix.length) }

        val reformattedReference = runInEdtAndGet { recommendationManager.reformatReference(originalReference, rangeMarker, invocationOffset) }
        assertThat(reformattedReference.licenseName()).isEqualTo("test_license")
        assertThat(reformattedReference.repository()).isEqualTo("test_repo")
        assertThat(reformattedReference.recommendationContentSpan().start()).isEqualTo(rangeMarker.startOffset - invocationOffset)
        assertThat(reformattedReference.recommendationContentSpan().end()).isEqualTo(rangeMarker.endOffset - invocationOffset - documentContentSurfix.length)
        val span = runInEdtAndGet {
            fixture.editor.document.charsSequence.subSequence(
                reformattedReference.recommendationContentSpan().start(),
                reformattedReference.recommendationContentSpan().end()
            )
        }
        // span should not include newline char
        assertThat(span.last()).isNotEqualTo('\n')
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.services.codewhisperer

import com.intellij.openapi.project.Project
import com.intellij.testFramework.DisposableRule
import com.intellij.testFramework.fixtures.CodeInsightTestFixture
import com.intellij.testFramework.runInEdtAndWait
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import software.amazon.awssdk.services.codewhispererruntime.model.Completion
import software.amazon.awssdk.services.codewhispererruntime.model.Reference
import software.amazon.awssdk.services.codewhispererruntime.model.Span
import software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow.CodeWhispererCodeReferenceManager
import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule

class CodeWhispererReferenceManagerTest {
    @Rule
    @JvmField
    var projectRule = PythonCodeInsightTestFixtureRule()

    @Rule
    @JvmField
    val disposableRule = DisposableRule()

    private val documentContentContent = "012345678\n9"
    private lateinit var fixture: CodeInsightTestFixture
    private lateinit var project: Project
    private val originalReference = Reference.builder()
        .licenseName("test_license")
        .repository("test_repo")
        .recommendationContentSpan(
            Span.builder().start(0).end(14).build()
        )
        .build()

    private val recommendation = Completion.builder()
        .references(originalReference)
        .content("test\nreference")
        .build()

    @Before
    fun setup() {
        fixture = projectRule.fixture
        project = projectRule.project

        fixture.configureByText("test.py", documentContentContent)
        runInEdtAndWait {
            fixture.editor.caretModel.moveToOffset(documentContentContent.length)
        }
    }

    @Test
    fun `test getReferenceLineNums return expected line numbers`() {
        val referenceManager = CodeWhispererCodeReferenceManager(project)
        assertThat(referenceManager.getReferenceLineNums(fixture.editor, 0, 1)).isEqualTo("1")
        assertThat(referenceManager.getReferenceLineNums(fixture.editor, 0, 10)).isEqualTo("1 to 2")
    }

    @Test
    fun `test getOriginalContent lines returns full reference lines`() {
        val referenceManager = CodeWhispererCodeReferenceManager(project)
        val expectedRecommendation = listOf<String>("test", "reference")
        assertThat(referenceManager.getOriginalContentLines(recommendation, 0)).isEqualTo(expectedRecommendation)
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.services.codewhisperer

import org.assertj.core.api.Assertions.assertThat
import org.junit.Test
import org.mockito.kotlin.any
import org.mockito.kotlin.doAnswer
import org.mockito.kotlin.stub
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsRequest
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsResponse
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.generateMockCompletionDetail
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.getReferenceInfo
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.metadata
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.sdkHttpResponse

class CodeWhispererReferencesTest : CodeWhispererTestBase() {

    @Test
    fun `test references with invalid content span will remove the references - index out of bound`() {
        testReferencesRangeValidity("test", 0, 5, invalid = true)
    }

    @Test
    fun `test references with invalid content span will remove the references - start greater than end`() {
        testReferencesRangeValidity("test", 2, 1, invalid = true)
    }

    @Test
    fun `test references with valid content span will not remove the references`() {
        testReferencesRangeValidity("test", 0, 4, invalid = false)
    }

    @Test
    fun `test references with valid(empty) content span will not remove the references`() {
        testReferencesRangeValidity("test", 2, 2, invalid = false)
    }

    private fun testReferencesRangeValidity(content: String, start: Int, end: Int, invalid: Boolean) {
        val referenceInfo = getReferenceInfo()
        val invalidReferencesResponse = GenerateCompletionsResponse.builder()
            .completions(
                generateMockCompletionDetail(content, referenceInfo.first, referenceInfo.second, start, end),
            )
            .nextToken("")
            .responseMetadata(metadata)
            .sdkHttpResponse(sdkHttpResponse)
            .build() as GenerateCompletionsResponse

        mockClient.stub {
            on {
                mockClient.generateCompletions(any<GenerateCompletionsRequest>())
            } doAnswer {
                invalidReferencesResponse
            }
        }

        withCodeWhispererServiceInvokedAndWait { states ->
            states.recommendationContext.details.forEach {
                assertThat(it.recommendation.references().isEmpty()).isEqualTo(invalid)
            }
            popupManagerSpy.popupComponents.acceptButton.doClick()
        }
    }
}


// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package software.aws.toolkits.jetbrains.services.codewhisperer

import org.assertj.core.api.Assertions.assertThat
import org.junit.Test
import org.mockito.kotlin.any
import org.mockito.kotlin.doAnswer
import org.mockito.kotlin.stub
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsRequest
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsResponse
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonFileName
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponse
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonTestLeftContext
import kotlin.random.Random

class CodeWhispererRightContextTest : CodeWhispererTestBase() {

    @Test
    fun `test recommendation equal to right context should not show recommendation`() {
        val rightContext = pythonResponse.completions()[0].content()
        setFileContext(pythonFileName, pythonTestLeftContext, rightContext)
        withCodeWhispererServiceInvokedAndWait { states ->
            val firstRecommendation = states.recommendationContext.details[0]
            assertThat(firstRecommendation.isDiscarded).isEqualTo(true)
        }
    }

    @Test
    fun `test right context resolution will remove reference span if reference is the same as right context`() {
        val rightContext = pythonResponse.completions()[0].content()
        setFileContext(pythonFileName, pythonTestLeftContext, rightContext)
        withCodeWhispererServiceInvokedAndWait { states ->
            val firstRecommendation = states.recommendationContext.details[0]
            assertThat(firstRecommendation.isDiscarded).isEqualTo(true)
            val details = states.recommendationContext.details
            details.forEach {
                assertThat(it.reformatted.references().isEmpty())
            }
        }
    }

    @Test
    fun `test right context resolution will update range of reference if reference overlaps with right context`() {
        val firstRecommendationContent = pythonResponse.completions()[0].content()
        val lastNewLineIndex = firstRecommendationContent.lastIndexOf('\n')
        val rightContext = firstRecommendationContent.substring(lastNewLineIndex)
        setFileContext(pythonFileName, pythonTestLeftContext, rightContext)
        withCodeWhispererServiceInvokedAndWait { states ->
            val firstRecommendation = states.recommendationContext.details[0]
            assertThat(firstRecommendation.isDiscarded).isEqualTo(false)
            val firstDetail = states.recommendationContext.details[0]
            val span = firstDetail.reformatted.references()[0].recommendationContentSpan()
            assertThat(span.start()).isEqualTo(0)
            assertThat(span.end()).isEqualTo(lastNewLineIndex)
        }
    }

    @Test
    fun `test right context resolution will not change reference range if reference does not overlap with right context`() {
        val firstRecommendationContent = pythonResponse.completions()[0].content()
        val lastNewLineIndex = firstRecommendationContent.lastIndexOf('\n')
        val rightContext = firstRecommendationContent.substring(lastNewLineIndex)
        setFileContext(pythonFileName, pythonTestLeftContext, rightContext)

        val referenceInfo = CodeWhispererTestUtil.getReferenceInfo()
        val referencesResponse = GenerateCompletionsResponse.builder()
            .completions(
                CodeWhispererTestUtil
                    .generateMockCompletionDetail(firstRecommendationContent, referenceInfo.first, referenceInfo.second, 0, lastNewLineIndex),
            )
            .nextToken("")
            .responseMetadata(CodeWhispererTestUtil.metadata)
            .sdkHttpResponse(CodeWhispererTestUtil.sdkHttpResponse)
            .build() as GenerateCompletionsResponse

        mockClient.stub {
            on {
                mockClient.generateCompletions(any<GenerateCompletionsRequest>())
            } doAnswer {
                referencesResponse
            }
        }

        withCodeWhispererServiceInvokedAndWait { states ->
            val firstRecommendation = states.recommendationContext.details[0]
            assertThat(firstRecommendation.isDiscarded).isEqualTo(false)
            val firstDetail = states.recommendationContext.details[0]
            val span = firstDetail.reformatted.references()[0].recommendationContentSpan()
            assertThat(span.start()).isEqualTo(0)
            assertThat(span.end()).isEqualTo(lastNewLineIndex)
        }
    }

    @Test
    fun `test suffix of recommendation equal to right context should truncate recommendation when remaining is single-line`() {
        val firstRecommendation = pythonResponse.completions()[0].content()
        val newLineIndex = firstRecommendation.indexOf('\n')
        val remainingLength = Random.nextInt(1, newLineIndex)
        val remaining = firstRecommendation.substring(0, remainingLength)
        val rightContext = pythonResponse.completions()[0].content().substring(remainingLength)
        setFileContext(pythonFileName, pythonTestLeftContext, rightContext)
        withCodeWhispererServiceInvokedAndWait { states ->
            assertThat(states.recommendationContext.details[0].reformatted.content()).isEqualTo(remaining)
            popupManagerSpy.popupComponents.acceptButton.doClick()
            assertThat(states.requestContext.editor.document.text).isEqualTo(pythonTestLeftContext + remaining + rightContext)
        }
    }

    @Test
    fun `test suffix of recommendation equal to right context should not truncate recommendation when remaining is multi-line`() {
        val firstRecommendation = pythonResponse.completions()[0].content()
        val newLineIndex = firstRecommendation.indexOf('\n')
        val remainingLength = Random.nextInt(newLineIndex, firstRecommendation.length)
        val rightContext = pythonResponse.completions()[0].content().substring(remainingLength)
        setFileContext(pythonFileName, pythonTestLeftContext, rightContext)
        withCodeWhispererServiceInvokedAndWait { states ->
            assertThat(states.recommendationContext.details[0].recommendation.content()).isEqualTo(firstRecommendation)
        }
    }

    @Test
    fun `test suffix of recommendation equal to prefix of right context should truncate recommendation when remaining is single-line`() {
        val firstRecommendation = pythonResponse.completions()[0].content()
        val newLineIndex = firstRecommendation.indexOf('\n')
        val remainingLength = Random.nextInt(1, newLineIndex)
        val remaining = firstRecommendation.substring(0, remainingLength)
        val rightContext = pythonResponse.completions()[0].content().substring(remainingLength) + "test"
        setFileContext(pythonFileName, pythonTestLeftContext, rightContext)
        withCodeWhispererServiceInvokedAndWait { states ->
            assertThat(states.recommendationContext.details[0].reformatted.content()).isEqualTo(remaining)
            popupManagerSpy.popupComponents.acceptButton.doClick()
            assertThat(states.requestContext.editor.document.text).isEqualTo(pythonTestLeftContext + remaining + rightContext)
        }
    }

    @Test
    fun `test suffix of recommendation equal to prefix of right context should not truncate recommendation when remaining is multi-line`() {
        val firstRecommendation = pythonResponse.completions()[0].content()
        val newLineIndex = firstRecommendation.indexOf('\n')
        val remainingLength = Random.nextInt(newLineIndex, firstRecommendation.length)
        val rightContext = pythonResponse.completions()[0].content().substring(remainingLength) + "test"
        setFileContext(pythonFileName, pythonTestLeftContext, rightContext)
        withCodeWhispererServiceInvokedAndWait { states ->
            assertThat(states.recommendationContext.details[0].recommendation.content()).isEqualTo(firstRecommendation)
        }
    }
}