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

package software.aws.toolkits.jetbrains.services.cloudwatch.logs

import com.intellij.ui.table.TableView
import com.intellij.util.ui.ListTableModel
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.channels.ClosedSendChannelException
import kotlinx.coroutines.runBlocking
import org.assertj.core.api.Assertions.assertThat
import org.assertj.core.api.Assertions.assertThatThrownBy
import org.junit.Before
import org.junit.Test
import org.mockito.Mockito
import org.mockito.kotlin.whenever
import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClient
import software.amazon.awssdk.services.cloudwatchlogs.model.GetLogEventsRequest
import software.amazon.awssdk.services.cloudwatchlogs.model.GetLogEventsResponse
import software.amazon.awssdk.services.cloudwatchlogs.model.OutputLogEvent
import software.amazon.awssdk.services.cloudwatchlogs.paginators.GetLogEventsIterable
import software.aws.toolkits.jetbrains.services.cloudwatch.logs.editor.LogStreamDateColumn
import software.aws.toolkits.jetbrains.services.cloudwatch.logs.editor.LogStreamMessageColumn
import software.aws.toolkits.jetbrains.utils.BaseCoroutineTest
import software.aws.toolkits.jetbrains.utils.waitForModelToBeAtLeast
import software.aws.toolkits.jetbrains.utils.waitForTrue
import software.aws.toolkits.resources.message
import java.time.Duration

@ExperimentalCoroutinesApi
class LogStreamListActorTest : BaseCoroutineTest() {
    private lateinit var client: CloudWatchLogsClient
    private lateinit var tableModel: ListTableModel<LogStreamEntry>
    private lateinit var table: TableView<LogStreamEntry>
    private lateinit var actor: CloudWatchLogsActor<LogStreamEntry>

    @Before
    fun loadVariables() {
        client = mockClientManagerRule.create()
        tableModel = ListTableModel(
            arrayOf(LogStreamDateColumn(), LogStreamMessageColumn()),
            mutableListOf<LogStreamEntry>()
        )
        table = TableView(tableModel)
        actor = LogStreamListActor(projectRule.project, client, table, "abc", "def")
    }

    @Test
    fun modelIsPopulated() {
        whenever(client.getLogEvents(Mockito.any<GetLogEventsRequest>()))
            .thenReturn(GetLogEventsResponse.builder().events(OutputLogEvent.builder().message("message").build()).build())
        runBlocking {
            actor.channel.send(CloudWatchLogsActor.Message.LoadInitial)
            tableModel.waitForModelToBeAtLeast(1)
            waitForTrue { table.emptyText.text == message("cloudwatch.logs.no_events") }
        }
        assertThat(tableModel.items.size).isOne()
        assertThat(tableModel.items.first().message).isEqualTo("message")
    }

    @Test
    fun modelIsPopulatedRange() {
        whenever(client.getLogEventsPaginator(Mockito.any<GetLogEventsRequest>()))
            .thenReturn(
                object : GetLogEventsIterable(client, null) {
                    override fun iterator() = mutableListOf(
                        GetLogEventsResponse.builder().events(
                            OutputLogEvent.builder().message("message").build()
                        ).build()
                    ).iterator()
                }
            )

        runBlocking {
            actor.channel.send(CloudWatchLogsActor.Message.LoadInitialRange(LogStreamEntry("@@@", 0), Duration.ofMillis(0)))
            tableModel.waitForModelToBeAtLeast(1)
            waitForTrue { table.emptyText.text == message("cloudwatch.logs.no_events") }
        }

        assertThat(tableModel.items.size).isOne()
        assertThat(tableModel.items.first().message).isEqualTo("message")
        assertThat(table.emptyText.text).isEqualTo(message("cloudwatch.logs.no_events"))
    }

    @Test
    fun emptyTableOnExceptionThrown() {
        whenever(client.getLogEvents(Mockito.any<GetLogEventsRequest>())).then { throw IllegalStateException("network broke") }
        runBlocking {
            actor.channel.send(CloudWatchLogsActor.Message.LoadInitial)
            waitForTrue {
                println(table.emptyText.text)
                table.emptyText.text == message("cloudwatch.logs.failed_to_load_stream", "def")
            }
        }
        assertThat(tableModel.items).isEmpty()
    }

    @Test
    fun emptyTableOnExceptionThrownRange() {
        whenever(client.getLogEvents(Mockito.any<GetLogEventsRequest>())).then { throw IllegalStateException("network broke") }
        runBlocking {
            actor.channel.send(CloudWatchLogsActor.Message.LoadInitialRange(LogStreamEntry("@@@", 0), Duration.ofMillis(0)))
            waitForTrue { table.emptyText.text == message("cloudwatch.logs.failed_to_load_stream", "def") }
        }
        assertThat(tableModel.items).isEmpty()
    }

    @Test
    fun loadingForwardAppendsToTable() {
        whenever(client.getLogEvents(Mockito.any<GetLogEventsRequest>()))
            .thenReturn(GetLogEventsResponse.builder().events(OutputLogEvent.builder().message("message").build()).nextForwardToken("2").build())
            .thenReturn(GetLogEventsResponse.builder().nextForwardToken("3").build())
            .thenReturn(GetLogEventsResponse.builder().events(OutputLogEvent.builder().message("message2").build()).nextForwardToken("4").build())
        runBlocking {
            actor.channel.send(CloudWatchLogsActor.Message.LoadInitial)
            tableModel.waitForModelToBeAtLeast(1)
        }
        assertThat(tableModel.items.size).isOne()
        assertThat(tableModel.items.first().message).isEqualTo("message")
        runBlocking {
            actor.channel.send(CloudWatchLogsActor.Message.LoadForward)
            actor.channel.send(CloudWatchLogsActor.Message.LoadForward)
            tableModel.waitForModelToBeAtLeast(2)
        }
        assertThat(tableModel.items.size).isEqualTo(2)
        assertThat(tableModel.items[1].message).isEqualTo("message2")
    }

    @Test
    fun loadingBackwardsPrependsToTable() {
        whenever(client.getLogEvents(Mockito.any<GetLogEventsRequest>()))
            .thenReturn(GetLogEventsResponse.builder().events(OutputLogEvent.builder().message("message").build()).nextBackwardToken("2").build())
            .thenReturn(GetLogEventsResponse.builder().nextBackwardToken("3").build())
            .thenReturn(
                GetLogEventsResponse.builder().events(OutputLogEvent.builder().message("message2").timestamp(3).build()).nextBackwardToken("2").build()
            )
        runBlocking {
            actor.channel.send(CloudWatchLogsActor.Message.LoadInitial)
            tableModel.waitForModelToBeAtLeast(1)
        }
        assertThat(tableModel.items.size).isOne()
        assertThat(tableModel.items.first().message).isEqualTo("message")
        runBlocking {
            actor.channel.send(CloudWatchLogsActor.Message.LoadBackward)
            actor.channel.send(CloudWatchLogsActor.Message.LoadBackward)
            tableModel.waitForModelToBeAtLeast(2)
        }
        assertThat(tableModel.items.size).isEqualTo(2)
        assertThat(tableModel.items.first().message).isEqualTo("message2")
        assertThat(tableModel.items.first().timestamp).isEqualTo(3)
        assertThat(tableModel.items[1].message).isEqualTo("message")
        assertThat(tableModel.items[1].timestamp).isZero()
    }

    @Test
    fun writeChannelAndCoroutineIsDisposed() {
        val channel = actor.channel
        actor.dispose()
        assertThatThrownBy {
            runBlocking {
                channel.send(CloudWatchLogsActor.Message.LoadBackward)
            }
        }.isInstanceOf(ClosedSendChannelException::class.java)
    }

    @Test
    fun loadInitialFilterThrows() {
        val client = mockClientManagerRule.create<CloudWatchLogsClient>()
        val tableModel = ListTableModel<LogStreamEntry>()
        val table = TableView(tableModel)
        val actor = LogStreamListActor(projectRule.project, client, table, "abc", "def")
        runBlocking {
            actor.channel.send(CloudWatchLogsActor.Message.LoadInitialFilter("abc"))
            waitForTrue { actor.channel.isClosedForSend }
        }
    }
}


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

package software.aws.toolkits.jetbrains.services.cloudwatch.logs

import com.intellij.openapi.application.runInEdt
import com.intellij.openapi.fileEditor.FileEditorManager
import com.intellij.testFramework.LightVirtualFile
import com.intellij.testFramework.ProjectRule
import com.intellij.testFramework.TestActionEvent
import com.intellij.ui.table.JBTable
import com.intellij.util.ui.ListTableModel
import kotlinx.coroutines.delay
import kotlinx.coroutines.runBlocking
import org.assertj.core.api.Assertions.assertThat
import org.junit.After
import org.junit.Rule
import org.junit.Test
import org.mockito.Mockito
import org.mockito.kotlin.whenever
import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClient
import software.amazon.awssdk.services.cloudwatchlogs.model.GetLogEventsRequest
import software.amazon.awssdk.services.cloudwatchlogs.model.GetLogEventsResponse
import software.amazon.awssdk.services.cloudwatchlogs.model.LogStream
import software.amazon.awssdk.services.cloudwatchlogs.model.OutputLogEvent
import software.amazon.awssdk.services.cloudwatchlogs.paginators.GetLogEventsIterable
import software.aws.toolkits.jetbrains.core.MockClientManagerRule
import software.aws.toolkits.jetbrains.services.cloudwatch.logs.actions.OpenCurrentInEditorAction
import software.aws.toolkits.jetbrains.services.cloudwatch.logs.actions.OpenLogStreamInEditorAction
import software.aws.toolkits.jetbrains.services.cloudwatch.logs.editor.LogStreamsStreamColumn

class OpenLogStreamInEditorActionTest {
    @JvmField
    @Rule
    val projectRule = ProjectRule()

    @JvmField
    @Rule
    val mockClientManagerRule = MockClientManagerRule()

    // TODO: figure out why this doesn't work on 223 Windows
//    @JvmField
//    @Rule
//    val timeout = CoroutinesTimeout.seconds(15)

    @After
    fun after() {
        val editorManager = FileEditorManager.getInstance(projectRule.project)
        runInEdt {
            editorManager.openFiles.forEach { editorManager.closeFile(it) }
        }
    }

    @Test
    fun testOpeningFileFromGroup() {
        val client = mockClientManagerRule.create<CloudWatchLogsClient>()
        whenever(client.getLogEventsPaginator(Mockito.any<GetLogEventsRequest>()))
            .thenReturn(
                object : GetLogEventsIterable(client, null) {
                    override fun iterator() = mutableListOf(
                        GetLogEventsResponse.builder().events(
                            OutputLogEvent.builder().message("abc").build(),
                            OutputLogEvent.builder().message("def").build()
                        ).build()
                    ).iterator()
                }
            )

        val tableModel = ListTableModel<LogStream>(LogStreamsStreamColumn(sortable = false))
        val table = JBTable(tableModel)
        tableModel.addRow(LogStream.builder().logStreamName("54321").build())
        // select the first row
        table.setRowSelectionInterval(0, 0)
        val action = OpenLogStreamInEditorAction(projectRule.project, client, "cloudWatchMock", "12345")
        action.actionPerformed(TestActionEvent())
        runBlocking {
            blockUntilFileOpen()
        }
        val openFile = FileEditorManager.getInstance(projectRule.project).openFiles.first()
        assertThat(openFile).isInstanceOf(LightVirtualFile::class.java)
        assertThat(openFile.isWritable).isFalse()
        assertThat((openFile as LightVirtualFile).content).isEqualTo("abc\ndef\n")
    }

    @Test
    fun testOpeningFileFromStream() {
        val tableModel = ListTableModel<LogStreamEntry>()
        tableModel.addRows(
            mutableListOf(
                LogStreamEntry("abc\n", 0L),
                LogStreamEntry("def\n", 0L),
                LogStreamEntry("ghi\n", 0L)
            )
        )
        val action = OpenCurrentInEditorAction(projectRule.project, "12345") { tableModel.items }
        action.actionPerformed(TestActionEvent())
        runBlocking {
            blockUntilFileOpen()
        }
        val openFile = FileEditorManager.getInstance(projectRule.project).openFiles.first()
        assertThat(openFile).isInstanceOf(LightVirtualFile::class.java)
        assertThat(openFile.isWritable).isFalse()
        assertThat((openFile as LightVirtualFile).content).isEqualTo("abc\ndef\nghi\n")
    }

    private suspend fun blockUntilFileOpen() {
        val editorManager = FileEditorManager.getInstance(projectRule.project)
        while (editorManager.openFiles.isEmpty()) {
            delay(20)
        }
    }
}


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

package software.aws.toolkits.jetbrains.services.cloudwatch.logs

import com.intellij.testFramework.ProjectRule
import com.intellij.testFramework.TestActionEvent
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.delay
import kotlinx.coroutines.runBlocking
import org.assertj.core.api.Assertions.assertThat
import org.junit.Rule
import org.junit.Test
import software.aws.toolkits.jetbrains.services.cloudwatch.logs.actions.TailLogsAction

class StreamLogsTest {
    @JvmField
    @Rule
    val projectRule = ProjectRule()

    // TODO: figure out why this doesn't work on 223 Windows
//    @JvmField
//    @Rule
//    val timeout = CoroutinesTimeout.seconds(10)

    @Test
    fun streamsWhenEnabled() {
        val channel = Channel<CloudWatchLogsActor.Message>()
        val tailLogs = TailLogsAction(projectRule.project) { channel }
        runBlocking {
            tailLogs.setSelected(TestActionEvent(), true)
            var response = channel.receive()
            assertThat(response).isInstanceOf(CloudWatchLogsActor.Message.LoadForward::class.java)
            response = channel.receive()
            assertThat(response).isInstanceOf(CloudWatchLogsActor.Message.LoadForward::class.java)
        }
    }

    @Test
    fun cancelsOnChannelClose() {
        val channel = Channel<CloudWatchLogsActor.Message>()
        val tailLogs = TailLogsAction(projectRule.project) { channel }
        channel.close()
        tailLogs.setSelected(TestActionEvent(), true)
        runBlocking {
            while (tailLogs.logStreamingJob?.isCompleted != true) {
                delay(10)
                println(tailLogs.logStreamingJob?.isActive)
            }
        }
        assertThat(tailLogs.logStreamingJob?.isActive).isFalse()
    }

    @Test
    fun cancelsOnCancel() {
        val channel = Channel<CloudWatchLogsActor.Message>()
        val tailLogs = TailLogsAction(projectRule.project) { channel }
        tailLogs.setSelected(TestActionEvent(), true)
        assertThat(tailLogs.logStreamingJob?.isActive).isTrue()
        tailLogs.setSelected(TestActionEvent(), false)
        runBlocking {
            while (tailLogs.logStreamingJob?.isCompleted != true) {
                delay(10)
            }
        }
        assertThat(tailLogs.logStreamingJob?.isActive).isFalse()
    }
}


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

package software.aws.toolkits.jetbrains.services.cloudwatch.logs

import com.intellij.ui.table.JBTable
import com.intellij.util.text.SyncDateFormat
import com.intellij.util.ui.ListTableModel
import org.assertj.core.api.Assertions.assertThat
import org.junit.Test
import org.mockito.kotlin.mock
import org.mockito.kotlin.whenever
import software.amazon.awssdk.services.cloudwatchlogs.model.LogStream
import software.aws.toolkits.jetbrains.services.cloudwatch.logs.editor.LogStreamsDateColumn
import software.aws.toolkits.jetbrains.services.cloudwatch.logs.editor.LogStreamsStreamColumn
import java.text.SimpleDateFormat
import java.time.ZoneId
import java.time.ZonedDateTime
import java.time.format.DateTimeFormatter
import java.util.Locale
import javax.swing.SortOrder

class TableUtilsTest {
    @Test
    fun `sorting must be enabled to work`() {
        assertThat(LogStreamsDateColumn(sortable = false).comparator).isNull()
        assertThat(LogStreamsDateColumn(sortable = true).comparator).isNotNull

        assertThat(LogStreamsStreamColumn(sortable = false).comparator).isNull()
        assertThat(LogStreamsStreamColumn(sortable = true).comparator).isNotNull
    }

    @Test
    fun `test short MDY sorting`() {
        val dateFormat = SimpleDateFormat("M/d/yy", Locale.ENGLISH)
        val table = JBTable(createTestMode(dateFormat))

        assertThat(table).satisfies {
            assertThat(it.getValueAt(0, 0)).isEqualTo("7/11/21")
            assertThat(it.getValueAt(1, 0)).isEqualTo("6/11/21")
            assertThat(it.getValueAt(2, 0)).isEqualTo("2/1/21")
        }
    }

    @Test
    fun `test medium MDY sorting`() {
        val dateFormat = SimpleDateFormat("MMM d, y", Locale.ENGLISH)
        val table = JBTable(createTestMode(dateFormat))

        assertThat(table).satisfies {
            assertThat(it.getValueAt(0, 0)).isEqualTo("Jul 11, 2021")
            assertThat(it.getValueAt(1, 0)).isEqualTo("Jun 11, 2021")
            assertThat(it.getValueAt(2, 0)).isEqualTo("Feb 1, 2021")
        }
    }

    @Test
    fun `test medium DMY sorting`() {
        val dateFormat = SimpleDateFormat("d MMM y", Locale.ENGLISH)
        val table = JBTable(createTestMode(dateFormat))

        assertThat(table).satisfies {
            assertThat(it.getValueAt(0, 0)).isEqualTo("11 Jul 2021")
            assertThat(it.getValueAt(1, 0)).isEqualTo("11 Jun 2021")
            assertThat(it.getValueAt(2, 0)).isEqualTo("1 Feb 2021")
        }
    }

    private fun createTestMode(format: SimpleDateFormat) = ListTableModel(
        arrayOf(LogStreamsDateColumn(sortable = true, SyncDateFormat(format))),
        mutableListOf(
            mockLogStream(stringToEpoch("2021-06-11T00:00:00-08:00")),
            mockLogStream(stringToEpoch("2021-02-01T00:00:00-08:00")),
            mockLogStream(stringToEpoch("2021-07-11T00:00:00-08:00")),
        ),
        0,
        SortOrder.DESCENDING
    )

    private fun stringToEpoch(string: String) = ZonedDateTime.parse(string, DateTimeFormatter.ISO_ZONED_DATE_TIME)
        .toLocalDateTime()
        .atZone(ZoneId.systemDefault())
        .toInstant()
        .toEpochMilli()

    private fun mockLogStream(epoch: Long): LogStream {
        val stream: LogStream = mock()
        whenever(stream.lastEventTimestamp()).thenReturn(epoch)

        return stream
    }
}


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

package software.aws.toolkits.jetbrains.services.cloudwatch.logs

import com.intellij.testFramework.ProjectRule
import com.intellij.testFramework.TestActionEvent
import com.intellij.ui.components.JBTextArea
import com.intellij.ui.table.TableView
import com.intellij.util.ui.ListTableModel
import org.assertj.core.api.Assertions.assertThat
import org.junit.Rule
import org.junit.Test
import org.mockito.kotlin.mock
import software.aws.toolkits.jetbrains.services.cloudwatch.logs.actions.WrapLogsAction
import software.aws.toolkits.jetbrains.services.cloudwatch.logs.editor.LogStreamDateColumn
import software.aws.toolkits.jetbrains.services.cloudwatch.logs.editor.LogStreamMessageColumn
import java.awt.Container

class WrapLogsActionTest {
    @JvmField
    @Rule
    val projectRule = ProjectRule()

    @Test
    fun wrappingChangesModel() {
        val model = ListTableModel<LogStreamEntry>(
            arrayOf(LogStreamDateColumn(), LogStreamMessageColumn()),
            listOf(mock())
        )
        val table = TableView(model)
        val wrapLogsAction = WrapLogsAction(projectRule.project) { table }
        wrapLogsAction.setSelected(TestActionEvent(), true)
        val wrappedComponent = model.columnInfos[1].getRenderer(null)!!.getTableCellRendererComponent(table, 0, false, false, 0, 0)

        val textArea = (wrappedComponent as Container).getComponent(0)

        assertThat(textArea).isInstanceOf(JBTextArea::class.java)
        assertThat((textArea as JBTextArea).lineWrap).isTrue()

        wrapLogsAction.setSelected(TestActionEvent(), false)
        val component = model.columnInfos[1].getRenderer(null)!!.getTableCellRendererComponent(table, 0, false, false, 0, 0)
        val textArea2 = (component as Container).getComponent(0)

        assertThat(textArea2).isInstanceOf(JBTextArea::class.java)
        assertThat((textArea2 as JBTextArea).lineWrap).isFalse()
    }
}


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

package software.aws.toolkits.jetbrains.services.cloudwatch.logs.insights

import com.intellij.testFramework.EdtRule
import com.intellij.testFramework.RuleChain
import kotlinx.coroutines.delay
import kotlinx.coroutines.runBlocking
import org.assertj.core.api.Assertions.assertThat
import org.assertj.core.api.Assertions.assertThatThrownBy
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import org.mockito.kotlin.any
import org.mockito.kotlin.whenever
import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClient
import software.amazon.awssdk.services.cloudwatchlogs.model.CloudWatchLogsException
import software.amazon.awssdk.services.cloudwatchlogs.model.GetLogRecordRequest
import software.amazon.awssdk.services.cloudwatchlogs.model.GetLogRecordResponse
import software.aws.toolkits.jetbrains.core.MockClientManagerRule
import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule
import software.aws.toolkits.jetbrains.utils.waitForModelToBeAtLeast

class DetailedLogRecordTest {
    val projectRule = JavaCodeInsightTestFixtureRule()

    @Rule
    @JvmField
    val ruleChain = RuleChain(projectRule, EdtRule())

    @JvmField
    @Rule
    val mockClientManagerRule = MockClientManagerRule()

    private lateinit var client: CloudWatchLogsClient

    @Before
    fun setUp() {
        client = mockClientManagerRule.create()
    }

    @Test
    fun `Loads record into table`() {
        whenever(client.getLogRecord(any<GetLogRecordRequest>()))
            .thenReturn(
                GetLogRecordResponse.builder().logRecord(mapOf("field" to "value", "message" to "message")).build()
            )
        val sut = DetailedLogRecord(projectRule.project, client, "ptr")
        val model = sut.tableView.listTableModel

        runBlocking {
            model.waitForModelToBeAtLeast(1)
        }

        assertThat(model.items).containsExactlyInAnyOrder("field" to "value", "message" to "message")
    }

    @Test
    fun `Does not throw on client exception`() {
        whenever(client.getLogRecord(any<GetLogRecordRequest>()))
            .thenThrow(CloudWatchLogsException::class.java)

        val sut = DetailedLogRecord(projectRule.project, client, "ptr")
        runBlocking {
            while (!sut.isLoaded()) {
                delay(10)
            }
        }

        assertThat(sut.tableView.listTableModel.items).isEmpty()
    }

    @Test
    fun getLogGroup() {
        assertThat(DetailedLogRecord.extractLogGroup("123456789012:/log/group/name")).isEqualTo("/log/group/name")
        assertThat(DetailedLogRecord.extractLogGroup("123456789012:1./group_with#symbols-name")).isEqualTo("1./group_with#symbols-name")

        assertThatThrownBy { DetailedLogRecord.extractLogGroup("123456789012:") }.isInstanceOf(IllegalStateException::class.java)
        assertThatThrownBy { DetailedLogRecord.extractLogGroup("/name") }.isInstanceOf(IllegalStateException::class.java)
        assertThatThrownBy { DetailedLogRecord.extractLogGroup("123:/name") }.isInstanceOf(IllegalStateException::class.java)
    }
}


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

package software.aws.toolkits.jetbrains.services.cloudwatch.logs.insights

import software.aws.toolkits.core.ConnectionSettings
import java.time.temporal.ChronoUnit

object InsightsUtils {
    fun queryDetails(
        connectionSettings: ConnectionSettings,
        logGroups: List<String> = listOf("logGroup"),
        timeRange: TimeRange = TimeRange.RelativeRange(1, ChronoUnit.DAYS),
        query: QueryString = QueryString.InsightsQueryString("query")
    ) = QueryDetails(
        connectionSettings = connectionSettings,
        logGroups = logGroups,
        timeRange = timeRange,
        query = query
    )
}


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

package software.aws.toolkits.jetbrains.services.cloudwatch.logs.insights

import com.intellij.ui.table.TableView
import com.intellij.util.ui.ListTableModel
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.delay
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.test.runBlockingTest
import org.assertj.core.api.Assertions.assertThat
import org.junit.After
import org.junit.Before
import org.junit.Test
import org.mockito.Mockito
import org.mockito.kotlin.any
import org.mockito.kotlin.verify
import org.mockito.kotlin.whenever
import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClient
import software.amazon.awssdk.services.cloudwatchlogs.model.CloudWatchLogsException
import software.amazon.awssdk.services.cloudwatchlogs.model.GetQueryResultsRequest
import software.amazon.awssdk.services.cloudwatchlogs.model.GetQueryResultsResponse
import software.amazon.awssdk.services.cloudwatchlogs.model.QueryStatus
import software.amazon.awssdk.services.cloudwatchlogs.model.ResultField
import software.amazon.awssdk.services.cloudwatchlogs.model.StopQueryRequest
import software.aws.toolkits.jetbrains.services.cloudwatch.logs.InsightsQueryResultsActor
import software.aws.toolkits.jetbrains.utils.BaseCoroutineTest
import software.aws.toolkits.jetbrains.utils.waitForModelToBeAtLeast
import software.aws.toolkits.jetbrains.utils.waitForTrue
import software.aws.toolkits.resources.message
import java.util.concurrent.CountDownLatch
import java.util.concurrent.TimeUnit

@ExperimentalCoroutinesApi
class QueryActorTest : BaseCoroutineTest() {
    private lateinit var client: CloudWatchLogsClient
    private lateinit var tableModel: ListTableModel<LogResult>
    private lateinit var table: TableView<LogResult>
    private lateinit var queryactor: InsightsQueryResultsActor

    @Before
    fun setUp() {
        client = mockClientManagerRule.create()
        tableModel = ListTableModel()
        table = TableView(tableModel)
        queryactor = InsightsQueryResultsActor(projectRule.project, client, table, "1234")
    }

    @After
    fun tearDown() {
        Mockito.reset(client)
    }

    @Test
    fun `dedupes events`() {
        val sampleResults = listOf(
            ResultField.builder()
                .field("@message")
                .value("First Sample Message")
                .build(),
            ResultField.builder()
                .field("@ptr")
                .value("ptr")
                .build()
        )

        whenever(client.getQueryResults(any<GetQueryResultsRequest>()))
            .thenReturn(
                GetQueryResultsResponse.builder().status(QueryStatus.RUNNING).results(sampleResults).build()
            )
            .thenReturn(
                GetQueryResultsResponse.builder().status(QueryStatus.COMPLETE).results(sampleResults).build()
            )

        runBlockingTest {
            queryactor.channel.send(InsightsQueryResultsActor.Message.StartLoadingAll)
            tableModel.waitForModelToBeAtLeast(1)
        }

        assertThat(tableModel.items.size).isOne()
        assertThat(tableModel.items.first().keys).contains("@message")
    }

    @Test
    fun `loads events`() {
        val sampleResult1 = ResultField.builder()
            .field("@ptr")
            .value("1234")
            .build()
        val firstSampleResultList = listOf(sampleResult1)
        val sampleResult2 = ResultField.builder()
            .field("@ptr")
            .value("5678")
            .build()
        val secondSampleResultList = listOf(sampleResult2)
        whenever(client.getQueryResults(Mockito.any<GetQueryResultsRequest>()))
            .thenReturn(
                GetQueryResultsResponse.builder().status(QueryStatus.RUNNING).results(firstSampleResultList).build()
            )
            .thenReturn(
                GetQueryResultsResponse.builder().status(QueryStatus.COMPLETE).results(firstSampleResultList, secondSampleResultList).build()
            )
        runBlockingTest {
            queryactor.channel.send(InsightsQueryResultsActor.Message.StartLoadingAll)
            tableModel.waitForModelToBeAtLeast(2)
        }
        assertThat(tableModel.items.size).isEqualTo(2)
        assertThat(tableModel.items[0].values).isEqualTo(setOf("1234"))
        assertThat(tableModel.items[1].values).isEqualTo(setOf("5678"))
    }

    @Test
    fun `loads partially`() {
        val sampleResults = listOf(
            ResultField.builder()
                .field("@message")
                .value("First Sample Message")
                .build(),
            ResultField.builder()
                .field("@ptr")
                .value("ptr")
                .build()
        )

        whenever(client.getQueryResults(any<GetQueryResultsRequest>()))
            .thenReturn(
                GetQueryResultsResponse.builder().status(QueryStatus.RUNNING).results(sampleResults).build()
            )
            .thenThrow(
                CloudWatchLogsException::class.java
            )

        runBlockingTest {
            queryactor.channel.send(InsightsQueryResultsActor.Message.StartLoadingAll)
            while (!queryactor.channel.isClosedForSend) {
                delay(10)
            }
        }

        assertThat(tableModel.items.size).isOne()
        assertThat(tableModel.items.first().keys).contains("@message")
    }

    @Test
    fun `no results`() {
        whenever(client.getQueryResults(Mockito.any<GetQueryResultsRequest>()))
            .thenReturn(
                GetQueryResultsResponse.builder().status(QueryStatus.COMPLETE).build()
            )
        runBlockingTest {
            queryactor.channel.send(InsightsQueryResultsActor.Message.StartLoadingAll)
            waitForTrue { table.emptyText.text == message("cloudwatch.logs.no_results_found") }
        }
        assertThat(tableModel.items).isEmpty()
    }

    @Test
    fun `errors immediately`() {
        whenever(client.getQueryResults(any<GetQueryResultsRequest>()))
            .thenThrow(
                CloudWatchLogsException::class.java
            )

        runBlocking {
            queryactor.channel.send(InsightsQueryResultsActor.Message.StartLoadingAll)
            waitForTrue { table.emptyText.text == message("cloudwatch.logs.query_results_table_error") }
        }

        assertThat(tableModel.items.size).isZero()
    }

    @Test
    fun `stop loading`() {
        val sampleResult1 = ResultField.builder()
            .field("@ptr")
            .value("p1234")
            .build()
        val firstSampleResultList = listOf(sampleResult1)
        val sampleResult2 = ResultField.builder()
            .field("@ptr")
            .value("5678")
            .build()
        val secondSampleResultList = listOf(sampleResult2)
        whenever(client.getQueryResults(any<GetQueryResultsRequest>()))
            .thenReturn(
                GetQueryResultsResponse.builder().status(QueryStatus.RUNNING).results(firstSampleResultList).build()
            )
            .thenAnswer {
                runBlocking {
                    delay(90_000)
                }
                GetQueryResultsResponse.builder().status(QueryStatus.COMPLETE).results(firstSampleResultList, secondSampleResultList).build()
            }

        val latch = CountDownLatch(1)
        whenever(client.stopQuery(any<StopQueryRequest>()))
            .thenAnswer {
                latch.countDown()
            }

        runBlockingTest {
            queryactor.channel.send(InsightsQueryResultsActor.Message.StartLoadingAll)
            tableModel.waitForModelToBeAtLeast(1)
            queryactor.channel.send(InsightsQueryResultsActor.Message.StopLoading)

            while (!queryactor.channel.isClosedForSend) {
                delay(10)
            }
        }

        assertThat(tableModel.items.size).isEqualTo(1)
        assertThat(tableModel.items[0].values).isEqualTo(setOf("p1234"))
        assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue()
        verify(client).stopQuery(StopQueryRequest.builder().queryId("1234").build())
    }
}


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

package software.aws.toolkits.jetbrains.services.cloudwatch.logs.insights

import com.intellij.testFramework.runInEdtAndGet
import kotlinx.coroutines.runBlocking
import org.assertj.core.api.Assertions.assertThat
import org.assertj.core.api.Assertions.within
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import org.mockito.kotlin.any
import org.mockito.kotlin.argumentCaptor
import org.mockito.kotlin.stub
import org.mockito.kotlin.verify
import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClient
import software.amazon.awssdk.services.cloudwatchlogs.model.LogGroup
import software.amazon.awssdk.services.cloudwatchlogs.model.StartQueryRequest
import software.amazon.awssdk.services.cloudwatchlogs.model.StartQueryResponse
import software.aws.toolkits.core.ConnectionSettings
import software.aws.toolkits.core.credentials.aToolkitCredentialsProvider
import software.aws.toolkits.core.region.anAwsRegion
import software.aws.toolkits.jetbrains.core.MockClientManagerRule
import software.aws.toolkits.jetbrains.core.MockResourceCacheRule
import software.aws.toolkits.jetbrains.services.cloudwatch.logs.insights.InsightsUtils.queryDetails
import software.aws.toolkits.jetbrains.services.cloudwatch.logs.resources.CloudWatchResources
import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule
import software.aws.toolkits.resources.message
import java.time.Duration
import java.time.Instant
import java.time.temporal.ChronoUnit
import java.util.Calendar
import java.util.Date

class QueryEditorDialogTest {
    @Rule
    @JvmField
    val projectRule = JavaCodeInsightTestFixtureRule()

    @JvmField
    @Rule
    val mockClientManagerRule = MockClientManagerRule()

    @JvmField
    @Rule
    val mockResourceCache = MockResourceCacheRule()

    private lateinit var view: QueryEditor
    private lateinit var sut: QueryEditorDialog
    private lateinit var client: CloudWatchLogsClient

    private val credentials = aToolkitCredentialsProvider()
    private val region = anAwsRegion()
    private val connectionSettings = ConnectionSettings(credentials, region)

    @Before
    fun setUp() {
        val project = projectRule.project
        client = mockClientManagerRule.create()
        mockResourceCache.addEntry(
            CloudWatchResources.LIST_LOG_GROUPS,
            region.id,
            credentials.id,
            listOf(LogGroup.builder().logGroupName("log1").build())
        )
        view = QueryEditor(
            project,
            queryDetails(
                connectionSettings = connectionSettings,
                logGroups = listOf()
            )
        )
        sut = runInEdtAndGet { QueryEditorDialog(project, connectionSettings, "log1") }
        runBlocking {
            // annoying race between view initialization and test assertion
            sut.setView(
                queryDetails(
                    connectionSettings = connectionSettings,
                    logGroups = listOf("log1")
                )
            )
        }

        client.stub {
            on(it.startQuery(any<StartQueryRequest>())).thenReturn(
                StartQueryResponse.builder().queryId("queryId").build()
            )
        }
    }

    @Test
    fun `Dialog selects correct log groups`() {
        mockResourceCache.addEntry(
            CloudWatchResources.LIST_LOG_GROUPS,
            region.id,
            credentials.id,
            listOf(
                LogGroup.builder().logGroupName("log0").build(),
                LogGroup.builder().logGroupName("log1").build(),
                LogGroup.builder().logGroupName("log2").build()
            )
        )
        assertThat(sut.getQueryDetails().logGroups).containsExactly("log1")
    }

    @Test
    fun `Dialog selects multiple log groups`() {
        mockResourceCache.addEntry(
            CloudWatchResources.LIST_LOG_GROUPS,
            region.id,
            credentials.id,
            listOf(
                LogGroup.builder().logGroupName("log0").build(),
                LogGroup.builder().logGroupName("log1").build(),
                LogGroup.builder().logGroupName("log2").build()
            )
        )
        runBlocking {
            val queryDetails = queryDetails(
                connectionSettings = connectionSettings,
                logGroups = listOf("log0", "log1")
            )
            sut = runInEdtAndGet { QueryEditorDialog(projectRule.project, queryDetails) }
            sut.setView(queryDetails)
        }
        assertThat(sut.getQueryDetails().logGroups).containsExactly("log0", "log1")
    }

    @Test
    fun `Dialog selects relative time correctly`() {
        runBlocking {
            val details = queryDetails(
                connectionSettings = connectionSettings,
                timeRange = TimeRange.RelativeRange(72761, ChronoUnit.DAYS)
            )
            sut.setView(details)
        }

        assertThat(sut.getQueryDetails().timeRange).isInstanceOfSatisfying(TimeRange.RelativeRange::class.java) {
            assertThat(it.relativeTimeAmount).isEqualTo(72761)
            assertThat(it.relativeTimeUnit).isEqualTo(ChronoUnit.DAYS)
        }
    }

    @Test
    fun `Dialog selects absolute time correctly`() {
        runBlocking {
            val details = queryDetails(
                connectionSettings = connectionSettings,
                timeRange = TimeRange.AbsoluteRange(
                    startDate = Date(1),
                    endDate = Date(1000000)
                )
            )
            sut.setView(details)
        }

        assertThat(sut.getQueryDetails().timeRange).isInstanceOfSatisfying(TimeRange.AbsoluteRange::class.java) {
            assertThat(it.startDate).isEqualTo(Date(1))
            assertThat(it.endDate).isEqualTo(Date(1000000))
        }
    }

    @Test
    fun `Dialog selects insights query box correctly`() {
        runBlocking {
            val details = queryDetails(
                connectionSettings = connectionSettings,
                query = QueryString.InsightsQueryString("insights query")
            )
            sut.setView(details)
        }

        assertThat(sut.getQueryDetails().query).isInstanceOfSatisfying(QueryString.InsightsQueryString::class.java) {
            assertThat(it.query).isEqualTo("insights query")
        }
    }

    @Test
    fun `Dialog selects search term box correctly`() {
        runBlocking {
            val details = queryDetails(
                connectionSettings = connectionSettings,
                query = QueryString.SearchTermQueryString("search term")

            )
            sut.setView(details)
        }

        assertThat(sut.getQueryDetails().query).isInstanceOfSatisfying(QueryString.SearchTermQueryString::class.java) {
            assertThat(it.searchTerm).isEqualTo("search term")
        }
    }

    @Test
    fun `validates absolute or relative time selected`() {
        setViewDetails(absoluteTime = false, relativeTime = false)
        assertThat(sut.validateEditorEntries(view)?.message).contains(message("cloudwatch.logs.validation.timerange"))
    }

    @Test
    fun `validates start date must be before end date`() {
        val cal = Calendar.getInstance()
        cal.add(Calendar.DATE, -1)
        setViewDetails(absoluteTime = true, startDate = Calendar.getInstance().time, endDate = cal.time)
        assertThat(sut.validateEditorEntries(view)?.message).contains(message("cloudwatch.logs.compare.start.end.date"))
    }

    @Test
    fun `validates unit specified for relative time`() {
        setViewDetails(relativeTime = true, relativeTimeNumber = "")
        assertThat(sut.validateEditorEntries(view)?.message).contains(message("cloudwatch.logs.no_relative_time_number"))
    }

    @Test
    fun `Search query type is always selected`() {
        setViewDetails(relativeTime = true)
        assertThat(!view.queryLogGroupsRadioButton.isSelected && !view.searchTerm.isSelected).isFalse()

        view.queryLogGroupsRadioButton.doClick()
        assertThat(view.queryLogGroupsRadioButton.isSelected).isTrue()
        assertThat(view.searchTerm.isSelected).isFalse()

        view.searchTerm.doClick()
        assertThat(view.queryLogGroupsRadioButton.isSelected).isFalse()
        assertThat(view.searchTerm.isSelected).isTrue()
    }

    @Test
    fun `validates search term entered`() {
        setViewDetails(relativeTime = true, querySearch = true, searchTerm = "")
        assertThat(sut.validateEditorEntries(view)?.message).contains(message("cloudwatch.logs.no_term_entered"))
    }

    @Test
    fun `validates query entered`() {
        setViewDetails(relativeTime = true, queryLogs = true, query = "")
        assertThat(sut.validateEditorEntries(view)?.message).contains(message("cloudwatch.logs.no_query_entered"))
    }

    @Test
    fun `validates log groups selected`() {
        setViewDetails(relativeTime = true, logGroups = emptyList())
        assertThat(sut.validateEditorEntries(view)?.message).contains(message("cloudwatch.logs.no_log_group"))
    }

    @Test
    fun `Path with relative time and queries correctly entered gets executed`() {
        setViewDetails(relativeTime = true, queryLogs = true, query = "fields @timestamp")
        assertThat(sut.validateEditorEntries(view)?.message).isNull()
    }

    @Test
    fun `Path with absolute time and a search term entered gets executed`() {
        val cal = Calendar.getInstance()
        cal.add(Calendar.DATE, -1)
        setViewDetails(absoluteTime = true, endDate = Calendar.getInstance().time, startDate = cal.time, querySearch = true, searchTerm = "Error")
        assertThat(sut.validateEditorEntries(view)?.message).isNull()
    }

    @Test
    fun `Fields extracted correctly from query string`() {
        val fieldsAsSecondPartOfQuery = "filter @message like /Error/ | fields @message"
        val noFieldsQuery = "filter @message like /Error/"
        val onlyFieldsQuery = "fields @logStream, @timestamp"
        val twoFieldsQuery = "fields @timestamp, @logStream | limit 10 | fields @message"
        val fieldsInFilterQuery = "filter @message like /fields/ | fields @logStream"
        assertThat(QueryEditorDialog.getFields(fieldsAsSecondPartOfQuery)).isEqualTo(listOf("@message"))
        assertThat(QueryEditorDialog.getFields(noFieldsQuery)).isEqualTo(listOf("@timestamp", "@message"))
        assertThat(QueryEditorDialog.getFields(onlyFieldsQuery)).isEqualTo(listOf("@logStream", "@timestamp"))
        assertThat(QueryEditorDialog.getFields(twoFieldsQuery)).isEqualTo(listOf("@timestamp", "@logStream", "@message"))
        assertThat(QueryEditorDialog.getFields(fieldsInFilterQuery)).isEqualTo(listOf("@logStream"))
    }

    @Test
    fun `startQuery with multiple log groups`() {
        val query = queryDetails(
            connectionSettings = connectionSettings,
            logGroups = mutableListOf("logGroup", "anotherLogGroup")
        )

        runBlocking { sut.startQueryAsync(query).await() }

        val captor = argumentCaptor<StartQueryRequest>()
        verify(client).startQuery(captor.capture())

        captor.firstValue.let {
            assertThat(it.logGroupNames()).containsExactly("logGroup", "anotherLogGroup")
        }
    }

    @Test
    fun `startQuery with absolute time range`() {
        val end = Instant.now()
        val start = end.minus(Duration.ofDays(1))
        val query = queryDetails(
            connectionSettings = connectionSettings,
            timeRange = TimeRange.AbsoluteRange(Date.from(start), Date.from(end))
        )

        runBlocking { sut.startQueryAsync(query).await() }

        val captor = argumentCaptor<StartQueryRequest>()
        verify(client).startQuery(captor.capture())

        captor.firstValue.let {
            assertThat(it.logGroupNames()).containsExactly("logGroup")
            assertThat(it.startTime()).isEqualTo(start.epochSecond)
            assertThat(it.endTime()).isEqualTo(end.epochSecond)
            assertThat(it.queryString()).isEqualTo("query")
        }
    }

    @Test
    fun `startQuery with relative time range`() {
        val end = Instant.now()
        val start = end.minus(Duration.ofDays(1))
        val query = queryDetails(
            connectionSettings = connectionSettings,
            timeRange = TimeRange.RelativeRange(1, ChronoUnit.DAYS)
        )

        runBlocking { sut.startQueryAsync(query).await() }

        val captor = argumentCaptor<StartQueryRequest>()
        verify(client).startQuery(captor.capture())

        captor.firstValue.let {
            assertThat(it.logGroupNames()).containsExactly("logGroup")
            // 5 second leeway since the actual timestamps are calculated at time of execution
            assertThat(it.startTime()).isCloseTo(start.epochSecond, within(5L))
            assertThat(it.endTime()).isCloseTo(end.epochSecond, within(5L))
            assertThat(it.queryString()).isEqualTo("query")
        }
    }

    @Test
    fun `startQuery with search term`() {
        val query = queryDetails(
            connectionSettings = connectionSettings,
            query = QueryString.SearchTermQueryString("query")
        )

        runBlocking { sut.startQueryAsync(query).await() }

        val captor = argumentCaptor<StartQueryRequest>()
        verify(client).startQuery(captor.capture())

        captor.firstValue.let {
            assertThat(it.logGroupNames()).containsExactly("logGroup")
            assertThat(it.queryString()).isEqualTo("fields @timestamp, @message | filter @message like /query/")
        }
    }

    @Test
    fun `startQuery with Insights query language query`() {
        val query = queryDetails(
            connectionSettings = connectionSettings,
            query = QueryString.InsightsQueryString("query")
        )

        runBlocking { sut.startQueryAsync(query).await() }

        val captor = argumentCaptor<StartQueryRequest>()
        verify(client).startQuery(captor.capture())

        captor.firstValue.let {
            assertThat(it.logGroupNames()).containsExactly("logGroup")
            assertThat(it.queryString()).isEqualTo("query")
        }
    }

    private fun setViewDetails(
        logGroups: List<String> = listOf("log1"),
        absoluteTime: Boolean = false,
        relativeTime: Boolean = false,
        startDate: Date = Calendar.getInstance().time,
        endDate: Date = Calendar.getInstance().time,
        relativeTimeUnit: String = "Minutes",
        relativeTimeNumber: String = "1",
        querySearch: Boolean = false,
        queryLogs: Boolean = false,
        searchTerm: String = "Example",
        query: String = "Example Query"
    ) {
        view.logGroupTable.populateLogGroups(logGroups.toSet(), logGroups)
        view.relativeTimeRadioButton.isSelected = relativeTime
        view.endDate.date = endDate
        view.startDate.date = startDate
        view.absoluteTimeRadioButton.isSelected = absoluteTime
        view.relativeTimeUnit.selectedItem = relativeTimeUnit
        view.relativeTimeNumber.text = relativeTimeNumber
        view.queryLogGroupsRadioButton.isSelected = queryLogs
        view.searchTerm.isSelected = querySearch
        view.querySearchTerm.text = searchTerm
        view.queryBox.text = query
    }
}


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

package software.aws.toolkits.jetbrains.services.cloudwatch.logs.insights

import com.intellij.testFramework.EdtRule
import com.intellij.testFramework.RuleChain
import com.intellij.testFramework.RunsInEdt
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import software.amazon.awssdk.services.cloudwatchlogs.model.QueryDefinition
import software.aws.toolkits.core.ConnectionSettings
import software.aws.toolkits.core.credentials.aToolkitCredentialsProvider
import software.aws.toolkits.core.region.anAwsRegion
import software.aws.toolkits.jetbrains.services.cloudwatch.logs.insights.InsightsUtils.queryDetails
import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule

@RunsInEdt
class RetrieveSavedQueryDialogTest {
    val projectRule = JavaCodeInsightTestFixtureRule()

    @Rule
    @JvmField
    val ruleChain = RuleChain(projectRule, EdtRule())

    private val credentials = aToolkitCredentialsProvider()
    private val region = anAwsRegion()
    private val connectionSettings = ConnectionSettings(credentials, region)

    // need to initialize under EDT with a valid application
    private lateinit var editor: QueryEditor

    @Before
    fun setUp() {
        editor = QueryEditor(
            projectRule.project,
            queryDetails(
                connectionSettings = connectionSettings
            )
        ).also {
            // initialize dummy values
            val groups = listOf("log0", "log1", "log2")
            it.logGroupTable.populateLogGroups(setOf("log1"), groups)
        }
    }

    @Test
    fun populateParentEditor_noLogGroups() {
        val definition = QueryDefinition.builder()
            .queryString("query")
            .build()
        RetrieveSavedQueryDialog.populateParentEditor(editor, definition)

        // no-ops log group selection
        assertThat(editor.logGroupTable.getSelectedLogGroups()).containsExactly("log1")
        assertThat(editor.queryLogGroupsRadioButton.isSelected).isTrue()
        assertThat(editor.queryBox.text).isEqualTo("query")
    }

    @Test
    fun populateParentEditor_withLogGroups() {
        val definition = QueryDefinition.builder()
            .logGroupNames("log0", "log1")
            .queryString("query")
            .build()
        RetrieveSavedQueryDialog.populateParentEditor(editor, definition)

        assertThat(editor.logGroupTable.getSelectedLogGroups()).containsExactly("log0", "log1")
        assertThat(editor.queryLogGroupsRadioButton.isSelected).isTrue()
        assertThat(editor.queryBox.text).isEqualTo("query")
    }
}


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

package software.aws.toolkits.jetbrains.services.cloudwatch.logs.insights

import com.intellij.testFramework.EdtRule
import com.intellij.testFramework.RuleChain
import com.intellij.testFramework.RunsInEdt
import kotlinx.coroutines.runBlocking
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import org.mockito.kotlin.argumentCaptor
import org.mockito.kotlin.doReturn
import org.mockito.kotlin.doThrow
import org.mockito.kotlin.stub
import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClient
import software.amazon.awssdk.services.cloudwatchlogs.model.CloudWatchLogsException
import software.amazon.awssdk.services.cloudwatchlogs.model.PutQueryDefinitionRequest
import software.amazon.awssdk.services.cloudwatchlogs.model.PutQueryDefinitionResponse
import software.amazon.awssdk.services.cloudwatchlogs.model.QueryDefinition
import software.aws.toolkits.core.ConnectionSettings
import software.aws.toolkits.core.credentials.aToolkitCredentialsProvider
import software.aws.toolkits.core.region.anAwsRegion
import software.aws.toolkits.jetbrains.core.MockClientManagerRule
import software.aws.toolkits.jetbrains.core.MockResourceCacheRule
import software.aws.toolkits.jetbrains.services.cloudwatch.logs.resources.CloudWatchResources
import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule
import software.aws.toolkits.resources.message

@RunsInEdt
class SaveQueryTest {
    val projectRule = JavaCodeInsightTestFixtureRule()

    @JvmField
    @Rule
    val ruleChain = RuleChain(projectRule, EdtRule())

    @JvmField
    @Rule
    val mockClientManagerRule = MockClientManagerRule()

    @JvmField
    @Rule
    val resourceCache = MockResourceCacheRule()

    private val credentials = aToolkitCredentialsProvider()
    private val region = anAwsRegion()
    private val connectionSettings = ConnectionSettings(credentials, region)
    private lateinit var client: CloudWatchLogsClient

    @Before
    fun setUp() {
        client = mockClientManagerRule.create()
    }

    @Test
    fun `Query name not entered, error message displayed`() {
        val project = projectRule.project
        val view = EnterQueryName()
        val validator = SaveQueryDialog(project, connectionSettings, "fields @timestamp", listOf("log1"))
        view.queryName.text = ""
        assertThat(validator.validateQueryName(view)?.message).contains(message("cloudwatch.logs.query_name_missing"))
    }

    @Test
    fun `Path with correctly entered Query name returns no validation error`() {
        val project = projectRule.project
        val view = EnterQueryName()
        val validator = SaveQueryDialog(project, connectionSettings, "fields @timestamp", listOf("log1"))
        view.queryName.text = "TrialQuery"
        assertThat(validator.validateQueryName(view)?.message).isNull()
    }

    @Test
    fun `Save new query`() {
        val putQueryDefinitionCaptor = argumentCaptor<PutQueryDefinitionRequest>()
        client.stub {
            on { putQueryDefinition(putQueryDefinitionCaptor.capture()) } doReturn PutQueryDefinitionResponse.builder().queryDefinitionId("1234").build()
        }
        resourceCache.addEntry(CloudWatchResources.DESCRIBE_QUERY_DEFINITIONS, region.id, credentials.id, listOf())

        val dialog = SaveQueryDialog(projectRule.project, connectionSettings, "fields @timestamp", listOf("log1"))
        dialog.view.queryName.text = "queryName"
        runBlocking {
            dialog.saveQuery().join()
        }

        assertThat(putQueryDefinitionCaptor.allValues).hasSize(1)
        assertThat(putQueryDefinitionCaptor.firstValue.queryDefinitionId()).isEqualTo(null)
        assertThat(putQueryDefinitionCaptor.firstValue.name()).isEqualTo("queryName")
        assertThat(putQueryDefinitionCaptor.firstValue.logGroupNames()).containsExactly("log1")
        assertThat(putQueryDefinitionCaptor.firstValue.queryString()).isEqualTo("fields @timestamp")
    }

    @Test
    fun `Save query when name already exists`() {
        val putQueryDefinitionCaptor = argumentCaptor<PutQueryDefinitionRequest>()
        val client = mockClientManagerRule.create<CloudWatchLogsClient>()
        val testQueryDefinition = QueryDefinition.builder().name("SampleQuery").queryDefinitionId("1234").build()
        client.stub {
            on { putQueryDefinition(putQueryDefinitionCaptor.capture()) } doReturn PutQueryDefinitionResponse.builder().queryDefinitionId("1234").build()
        }
        resourceCache.addEntry(CloudWatchResources.DESCRIBE_QUERY_DEFINITIONS, region.id, credentials.id, listOf(testQueryDefinition))

        val dialog = SaveQueryDialog(projectRule.project, connectionSettings, "fields @timestamp", listOf("log1"))
        dialog.view.queryName.text = "SampleQuery"
        runBlocking {
            dialog.saveQuery().join()
        }

        assertThat(putQueryDefinitionCaptor.allValues).hasSize(1)
        assertThat(putQueryDefinitionCaptor.firstValue.queryDefinitionId()).isEqualTo("1234")
        assertThat(putQueryDefinitionCaptor.firstValue.name()).isEqualTo("SampleQuery")
        assertThat(putQueryDefinitionCaptor.firstValue.logGroupNames()).containsExactly("log1")
        assertThat(putQueryDefinitionCaptor.firstValue.queryString()).isEqualTo("fields @timestamp")
    }

    @Test
    fun `Save query doesn't throw on failure`() {
        val putQueryDefinitionCaptor = argumentCaptor<PutQueryDefinitionRequest>()
        val client = mockClientManagerRule.create<CloudWatchLogsClient>()
        val testQueryDefinition = QueryDefinition.builder().name("SampleQuery").queryDefinitionId("1234").build()
        client.stub {
            on { putQueryDefinition(putQueryDefinitionCaptor.capture()) } doThrow CloudWatchLogsException::class
        }
        resourceCache.addEntry(CloudWatchResources.DESCRIBE_QUERY_DEFINITIONS, region.id, credentials.id, listOf(testQueryDefinition))

        val dialog = SaveQueryDialog(projectRule.project, connectionSettings, "fields @timestamp", listOf("log1"))
        dialog.view.queryName.text = "SampleQuery"
        runBlocking {
            dialog.saveQuery().join()
        }
    }
}


// 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 org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Test
import software.aws.toolkits.jetbrains.services.codewhisperer.util.BM25
import software.aws.toolkits.jetbrains.services.codewhisperer.util.BM250kapi
import software.aws.toolkits.jetbrains.services.codewhisperer.util.BM25Result

class Bm25Test {
    private lateinit var sut: BM25
    private lateinit var tokenizer: (String) -> List<String>

    @Before
    fun setup() {
        tokenizer = { str: String ->
            str.split(" ")
        }
    }

    @Test
    fun `scores should return all scores and topN should return only highest n - simple case 1`() {
        val query = "windy London"
        val corpus = listOf(
            "Hello there good man!",
            "It is quite windy in London",
            "How is the weather today?"
        )

        sut = BM250kapi(corpus, tokenizer)
        val scores = sut.score(query)

        assertThat(scores).isEqualTo(
            listOf(
                BM25Result(docString = "Hello there good man!", score = 0.0),
                BM25Result(docString = "It is quite windy in London", score = 0.937294722506405),
                BM25Result(docString = "How is the weather today?", score = 0.0),
            )
        )

        assertThat(sut.topN(query, 1)).isEqualTo(
            listOf(
                BM25Result(docString = "It is quite windy in London", score = 0.937294722506405)
            )
        )
    }

    @Test
    fun `scores should return correct pair of document to its calculated score - simple case 2`() {
        val query = "codewhisperer is a machine learning powered code generator"
        val corpus = listOf(
            "codewhisperer goes GA at April 2023",
            "machine learning tool is the trending topic!!! :)",
            "codewhisperer is good =))))",
            "codewhisperer vs. copilot, which code generator better?",
            "copilot is a AI code generator too",
            "it is so amazing!!"
        )

        sut = BM250kapi(corpus, tokenizer)
        val scores = sut.score(query)

        assertThat(scores).isEqualTo(
            listOf(
                BM25Result(docString = "codewhisperer goes GA at April 2023", score = 0.0),
                BM25Result(docString = "machine learning tool is the trending topic!!! :)", score = 2.5075225906025422),
                BM25Result(docString = "codewhisperer is good =))))", score = 0.33539413050870837),
                BM25Result(docString = "codewhisperer vs. copilot, which code generator better?", score = 1.0935565858644076),
                BM25Result(docString = "copilot is a AI code generator too", score = 2.5673872789459544),
                BM25Result(docString = "it is so amazing!!", score = 0.33539413050870837)
            )
        )

        assertThat(sut.topN(query, 1)).isEqualTo(
            listOf(
                BM25Result(docString = "copilot is a AI code generator too", score = 2.5673872789459544)
            )
        )
        assertThat(sut.topN(query, 3)).isEqualTo(
            listOf(
                BM25Result(docString = "copilot is a AI code generator too", score = 2.5673872789459544),
                BM25Result(docString = "machine learning tool is the trending topic!!! :)", score = 2.5075225906025422),
                BM25Result(docString = "codewhisperer vs. copilot, which code generator better?", score = 1.0935565858644076),
            )
        )
    }
}


// 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.ide.DataManager
import com.intellij.openapi.actionSystem.ActionManager
import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.actionSystem.DataContext
import com.intellij.openapi.actionSystem.IdeActions
import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_MOVE_CARET_LEFT
import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_MOVE_CARET_RIGHT
import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_TAB
import com.intellij.openapi.editor.actionSystem.EditorActionManager
import com.intellij.testFramework.runInEdtAndWait
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
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.javaFileName
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.javaResponse
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.javaTestContext
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.metadata
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.sdkHttpResponse
import software.aws.toolkits.jetbrains.services.codewhisperer.actions.CodeWhispererActionPromoter
import kotlin.test.fail

class CodeWhispererAcceptTest : CodeWhispererTestBase() {

    @Before
    override fun setUp() {
        super.setUp()

        // Use java code to test curly braces behavior
        projectRule.fixture.configureByText(javaFileName, javaTestContext)
        mockClient.stub {
            on {
                mockClient.generateCompletions(any<GenerateCompletionsRequest>())
            } doAnswer {
                javaResponse
            }
        }
        runInEdtAndWait {
            projectRule.fixture.editor.caretModel.moveToOffset(projectRule.fixture.editor.document.textLength - 2)
        }
    }

    @Test
    fun `test accept recommendation with no typeahead no matching brackets on the right`() {
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "")
    }

    @Test
    fun `test accept recommendation with no typeahead no matching brackets on the right using keyboard`() {
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "", useKeyboard = true)
    }

    @Test
    fun `test accept recommendation with no typeahead with matching brackets on the right`() {
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "(")
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "()")
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "() {")
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "{")
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "(){")
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "() {")
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "() {\n    }")
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "() {\n    }     ")
    }

    @Test
    fun `test accept recommendation with typeahead with matching brackets on the right`() {
        val lastIndexOfNewLine = javaResponse.completions()[0].content().lastIndexOf("\n")
        val recommendation = javaResponse.completions()[0].content().substring(0, lastIndexOfNewLine)
        recommendation.indices.forEach {
            val typeahead = recommendation.substring(0, it + 1)
            testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight(typeahead, "(")
            testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight(typeahead, "()")
            testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight(typeahead, "() {")
            testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight(typeahead, "{")
            testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight(typeahead, "(){")
            testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight(typeahead, "() {")
            testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight(typeahead, "() {\n    }")
            testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight(typeahead, "() {\n    }     ")
        }
    }

    @Test
    fun `test accept single-line recommendation with no typeahead with partial matching brackets on the right`() {
        mockClient.stub {
            on {
                mockClient.generateCompletions(any<GenerateCompletionsRequest>())
            } doAnswer {
                GenerateCompletionsResponse.builder()
                    .completions(
                        generateMockCompletionDetail("(x, y) {"),
                    )
                    .nextToken("")
                    .responseMetadata(metadata)
                    .sdkHttpResponse(sdkHttpResponse)
                    .build() as GenerateCompletionsResponse
            }
        }

        // any non-matching first-line right context should remain
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "(", "test")
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "()", "test")
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "() {", "test")
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "({", "test")
    }

    @Test
    fun `test accept multi-line recommendation with no typeahead with partial matching brackets on the right`() {
        // any first-line right context should be overwritten
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "(test")
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "()test")
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "() {test")
        testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight("", "({test")
    }

    @Test
    fun `test CodeWhisperer keyboard shortcuts should be prioritized to be executed`() {
        testCodeWhispererKeyboardShortcutShouldBePrioritized(ACTION_EDITOR_TAB)
        testCodeWhispererKeyboardShortcutShouldBePrioritized(ACTION_EDITOR_MOVE_CARET_RIGHT)
        testCodeWhispererKeyboardShortcutShouldBePrioritized(ACTION_EDITOR_MOVE_CARET_LEFT)
    }

    private fun testCodeWhispererKeyboardShortcutShouldBePrioritized(actionId: String) {
        projectRule.fixture.configureByText(javaFileName, javaTestContext)
        val wrongAction = object : AnAction() {
            override fun actionPerformed(e: AnActionEvent) {
                fail("the other action should not get executed first")
            }
        }
        withCodeWhispererServiceInvokedAndWait {
            val codeWhispererAction = ActionManager.getInstance().getAction(actionId)
            val actions = listOf<AnAction>(wrongAction, codeWhispererAction)
            val newActions = CodeWhispererActionPromoter().promote(
                actions.toMutableList(),
                DataManager.getInstance().getDataContext(projectRule.fixture.editor.contentComponent)
            )
            assertThat(newActions[0]).isEqualTo(codeWhispererAction)
            popupManagerSpy.popupComponents.acceptButton.doClick()
        }
    }

    private fun testAcceptRecommendationWithTypingAndMatchingBracketsOnTheRight(
        typing: String,
        brackets: String,
        remaining: String = "",
        useKeyboard: Boolean = false
    ) {
        projectRule.fixture.configureByText(javaFileName, buildContextWithRecommendation(brackets + remaining))
        runInEdtAndWait {
            // move the cursor to the correct trigger point (...void main<trigger>)
            projectRule.fixture.editor.caretModel.moveToOffset(47)
        }
        withCodeWhispererServiceInvokedAndWait { states ->
            val recommendation = states.recommendationContext.details[0].reformatted.content()
            val editor = projectRule.fixture.editor
            val expectedContext = buildContextWithRecommendation(recommendation + remaining)
            val startOffset = editor.caretModel.offset
            typing.forEachIndexed { index, char ->
                if (index < editor.caretModel.offset - startOffset) return@forEachIndexed
                projectRule.fixture.type(char)
            }
            acceptHelper(useKeyboard)
            assertThat(editor.document.text).isEqualTo(expectedContext)
        }
    }

    private fun acceptHelper(useKeyboard: Boolean) {
        if (useKeyboard) {
            EditorActionManager.getInstance().getActionHandler(IdeActions.ACTION_EDITOR_TAB)
                .execute(projectRule.fixture.editor, null, DataContext.EMPTY_CONTEXT)
        } else {
            popupManagerSpy.popupComponents.acceptButton.doClick()
        }
    }

    private fun buildContextWithRecommendation(recommendation: String): String {
        val lastIndexOfNewLine = javaTestContext.length - 2
        return javaTestContext.substring(0, lastIndexOfNewLine) + recommendation + javaTestContext.substring(lastIndexOfNewLine)
    }
}


// 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.junit.Before
import org.junit.Rule
import org.junit.Test
import org.mockito.kotlin.mock
import org.mockito.kotlin.spy
import org.mockito.kotlin.times
import org.mockito.kotlin.verify
import org.mockito.kotlin.whenever
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanManager
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.nodes.CodeWhispererActionNode
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.nodes.OpenCodeReferenceNode
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.nodes.RunCodeScanNode
import software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow.CodeWhispererCodeReferenceManager

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

    @JvmField
    @Rule
    val projectRule = ProjectRule()

    @JvmField
    @Rule
    val disposableRule = DisposableRule()

    private lateinit var project: Project
    private lateinit var sut: CodeWhispererActionNode
    private lateinit var explorerManager: CodeWhispererExplorerActionManager
    private lateinit var connectionManager: ToolkitConnectionManager
    private lateinit var codeScanManager: CodeWhispererCodeScanManager

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

        explorerManager = spy()
        ApplicationManager.getApplication().replaceService(CodeWhispererExplorerActionManager::class.java, explorerManager, disposableRule.disposable)

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

        codeScanManager = mock()
        project.replaceService(CodeWhispererCodeScanManager::class.java, codeScanManager, disposableRule.disposable)
    }

    @Test
    fun `openCodeReferenceNode`() {
        sut = OpenCodeReferenceNode(project)
        val referenceManager: CodeWhispererCodeReferenceManager = mock()
        project.replaceService(CodeWhispererCodeReferenceManager::class.java, referenceManager, disposableRule.disposable)

        sut.onDoubleClick(mock())

        verify(referenceManager).showCodeReferencePanel()
    }

    @Test
    fun `runCodeScanNode`() {
        whenever(codeScanManager.getActionButtonIconForExplorerNode()).thenReturn(mock())
        whenever(codeScanManager.getActionButtonText()).thenReturn("")

        sut = RunCodeScanNode(project)

        sut.onDoubleClick(mock())

        verify(codeScanManager).runCodeScan()
    }
}


// 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.ide.browsers.BrowserLauncher
import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.actionSystem.CommonDataKeys
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.options.Configurable
import com.intellij.openapi.options.ShowSettingsUtil
import com.intellij.testFramework.TestActionEvent
import com.intellij.testFramework.replaceService
import com.intellij.testFramework.runInEdtAndWait
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Test
import org.mockito.Mockito.mock
import org.mockito.Mockito.spy
import org.mockito.Mockito.verify
import org.mockito.kotlin.any
import org.mockito.kotlin.argumentCaptor
import org.mockito.kotlin.doNothing
import org.mockito.kotlin.eq
import software.aws.toolkits.jetbrains.services.codewhisperer.actions.CodeWhispererLearnMoreAction
import software.aws.toolkits.jetbrains.services.codewhisperer.actions.CodeWhispererLoginLearnMoreAction
import software.aws.toolkits.jetbrains.services.codewhisperer.actions.CodeWhispererShowSettingsAction
import software.aws.toolkits.jetbrains.services.codewhisperer.actions.CodeWhispererSsoLearnMoreAction
import software.aws.toolkits.jetbrains.services.codewhisperer.actions.CodeWhispererWhatIsAction
import software.aws.toolkits.jetbrains.services.codewhisperer.actions.DoNotShowAgainActionError
import software.aws.toolkits.jetbrains.services.codewhisperer.actions.DoNotShowAgainActionWarn
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererConfigurable
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
import java.net.URI

class CodeWhispererActionTest : CodeWhispererTestBase() {
    private lateinit var browserLauncherSpy: BrowserLauncher
    private lateinit var event: AnActionEvent

    @Before
    override fun setUp() {
        super.setUp()
        browserLauncherSpy = mock(BrowserLauncher::class.java)
        event = TestActionEvent { key ->
            when (key) {
                CommonDataKeys.PROJECT.name -> projectRule.project
                else -> null
            }
        }
        ApplicationManager.getApplication().replaceService(BrowserLauncher::class.java, browserLauncherSpy, disposableRule.disposable)
    }

    @Test
    fun `CodeWhispererLearnMoreAction actionPerformed should open correct URI`() {
        val action = CodeWhispererLearnMoreAction()
        testBrowserActionHelper(action, CodeWhispererConstants.CODEWHISPERER_LEARN_MORE_URI)
    }

    @Test
    fun `CodeWhispererSsoLearnMoreAction actionPerformed should open correct URI`() {
        val action = CodeWhispererSsoLearnMoreAction()
        testBrowserActionHelper(action, CodeWhispererConstants.CODEWHISPERER_SSO_LEARN_MORE_URI)
    }

    @Test
    fun `CodeWhispererLoginLearnMoreAction actionPerformed should open correct URI`() {
        val action = CodeWhispererLoginLearnMoreAction()
        testBrowserActionHelper(action, CodeWhispererConstants.CODEWHISPERER_LOGIN_LEARN_MORE_URI)
    }

    @Test
    fun `CodeWhispererWhatIsAction update should be enabled and visible when terms of service have been accepted`() {
        val action = CodeWhispererWhatIsAction()
        mockCodeWhispererEnabledStatus(true)
        action.update(event)
        assertThat(event.presentation.isEnabledAndVisible).isTrue
    }

    @Test
    fun `CodeWhispererWhatIsAction update should be disabled and invisible when terms of service have not been accepted`() {
        val action = CodeWhispererWhatIsAction()
        mockCodeWhispererEnabledStatus(false)
        action.update(event)
        assertThat(event.presentation.isEnabledAndVisible).isFalse
    }

    @Test
    fun `CodeWhispererWhatIsAction actionPerformed should open correct URI`() {
        val action = CodeWhispererWhatIsAction()
        testBrowserActionHelper(action, CodeWhispererConstants.CODEWHISPERER_LEARN_MORE_URI)
    }

    @Test
    fun `DoNotShowAgainActionWarn actionPerformed should set doNotShowAgainWarn`() {
        CodeWhispererExplorerActionManager.getInstance().setDoNotShowAgainWarn(false)
        val action = DoNotShowAgainActionWarn()
        action.actionPerformed(event)
        assertThat(CodeWhispererExplorerActionManager.getInstance().getDoNotShowAgainWarn()).isTrue
    }

    @Test
    fun `DoNotShowAgainActionError actionPerformed should set doNotShowAgainError`() {
        CodeWhispererExplorerActionManager.getInstance().setDoNotShowAgainError(false)
        val action = DoNotShowAgainActionError()
        action.actionPerformed(event)
        assertThat(CodeWhispererExplorerActionManager.getInstance().getDoNotShowAgainError()).isTrue
    }

    @Test
    fun `CodeWhispererShowSettingsAction actionPerformed should show settings dialog`() {
        val settingsSpy = spy(ShowSettingsUtil.getInstance())
        doNothing().`when`(settingsSpy).showSettingsDialog(any(), any<Class<out Configurable>>())
        ApplicationManager.getApplication().replaceService(ShowSettingsUtil::class.java, settingsSpy, disposableRule.disposable)
        val action = CodeWhispererShowSettingsAction()
        runInEdtAndWait {
            action.actionPerformed(event)
            verify(settingsSpy).showSettingsDialog(any(), eq(CodeWhispererConfigurable::class.java))
        }
    }

    private fun testBrowserActionHelper(action: AnAction, expectedUri: String) {
        val urlCaptor = argumentCaptor<URI>()
        action.actionPerformed(event)
        verify(browserLauncherSpy).browse(urlCaptor.capture())
        assertThat(urlCaptor.lastValue).isEqualTo(URI(expectedUri))
    }
}


// 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.testFramework.DisposableRule
import org.junit.Rule
import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule

open class CodeWhispererBasicTestBase {
    @Rule
    @JvmField
    val projectRule = PythonCodeInsightTestFixtureRule()

    @Rule
    @JvmField
    val disposableRule = DisposableRule()
}


// 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.testFramework.DisposableRule
import com.intellij.testFramework.ProjectRule
import com.intellij.testFramework.RuleChain
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.argumentCaptor
import org.mockito.kotlin.doReturn
import org.mockito.kotlin.doReturnConsecutively
import org.mockito.kotlin.mock
import org.mockito.kotlin.stub
import org.mockito.kotlin.times
import org.mockito.kotlin.verify
import org.mockito.kotlin.verifyNoInteractions
import org.mockito.kotlin.whenever
import software.amazon.awssdk.services.codewhisperer.CodeWhispererClient
import software.amazon.awssdk.services.codewhisperer.model.ArtifactType
import software.amazon.awssdk.services.codewhisperer.model.CodeScanFindingsSchema
import software.amazon.awssdk.services.codewhisperer.model.CreateCodeScanRequest
import software.amazon.awssdk.services.codewhisperer.model.CreateCodeScanResponse
import software.amazon.awssdk.services.codewhisperer.model.CreateCodeScanUploadUrlRequest
import software.amazon.awssdk.services.codewhisperer.model.CreateCodeScanUploadUrlResponse
import software.amazon.awssdk.services.codewhisperer.model.GetCodeScanRequest
import software.amazon.awssdk.services.codewhisperer.model.GetCodeScanResponse
import software.amazon.awssdk.services.codewhisperer.model.ListCodeScanFindingsRequest
import software.amazon.awssdk.services.codewhisperer.model.ListCodeScanFindingsResponse
import software.amazon.awssdk.services.codewhisperer.model.ProgrammingLanguage
import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
import software.amazon.awssdk.services.codewhispererruntime.model.CodeAnalysisStatus
import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlRequest
import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlResponse
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsRequest
import software.amazon.awssdk.services.codewhispererruntime.model.GetCodeAnalysisRequest
import software.amazon.awssdk.services.codewhispererruntime.model.GetCodeAnalysisResponse
import software.amazon.awssdk.services.codewhispererruntime.model.ListCodeAnalysisFindingsRequest
import software.amazon.awssdk.services.codewhispererruntime.model.ListCodeAnalysisFindingsResponse
import software.amazon.awssdk.services.codewhispererruntime.model.StartCodeAnalysisRequest
import software.amazon.awssdk.services.codewhispererruntime.model.StartCodeAnalysisResponse
import software.amazon.awssdk.services.codewhispererruntime.paginators.GenerateCompletionsIterable
import software.amazon.awssdk.services.ssooidc.SsoOidcClient
import software.aws.toolkits.core.TokenConnectionSettings
import software.aws.toolkits.core.utils.test.aString
import software.aws.toolkits.jetbrains.core.MockClientManagerRule
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
import software.aws.toolkits.jetbrains.core.credentials.BearerSsoConnection
import software.aws.toolkits.jetbrains.core.credentials.ManagedSsoProfile
import software.aws.toolkits.jetbrains.core.credentials.MockCredentialManagerRule
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_REGION
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.metadata
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonRequest
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponseWithToken
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.sdkHttpResponse
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptorImpl

class CodeWhispererClientAdaptorTest {
    val projectRule = ProjectRule()
    val disposableRule = DisposableRule()
    val mockClientManagerRule = MockClientManagerRule()
    val mockCredentialRule = MockCredentialManagerRule()
    val authManagerRule = MockToolkitAuthManagerRule()

    @Rule
    @JvmField
    val ruleChain = RuleChain(projectRule, mockCredentialRule, mockClientManagerRule, disposableRule)

    private lateinit var sigv4Client: CodeWhispererClient
    private lateinit var bearerClient: CodeWhispererRuntimeClient
    private lateinit var ssoClient: SsoOidcClient

    private lateinit var sut: CodeWhispererClientAdaptor
    private lateinit var connectionManager: ToolkitConnectionManager

    @Before
    fun setup() {
        sut = CodeWhispererClientAdaptorImpl(projectRule.project)
        ssoClient = mockClientManagerRule.create()

        sigv4Client = mockClientManagerRule.create<CodeWhispererClient>().stub {
            on { createCodeScanUploadUrl(any<CreateCodeScanUploadUrlRequest>()) } doReturn createCodeScanUploadUrlResponse
            on { createCodeScan(any<CreateCodeScanRequest>()) } doReturn createCodeScanResponse
            on { getCodeScan(any<GetCodeScanRequest>()) } doReturn getCodeScanResponse
            on { listCodeScanFindings(any<ListCodeScanFindingsRequest>()) } doReturn listCodeScanFindingsResponse
        }

        bearerClient = mockClientManagerRule.create<CodeWhispererRuntimeClient>().stub {
            on { generateCompletionsPaginator(any<GenerateCompletionsRequest>()) } doReturn generateCompletionsPaginatorResponse
            on { createUploadUrl(any<CreateUploadUrlRequest>()) } doReturn createUploadUrlResponse
            on { startCodeAnalysis(any<StartCodeAnalysisRequest>()) } doReturn startCodeAnalysisResponse
            on { getCodeAnalysis(any<GetCodeAnalysisRequest>()) } doReturn getCodeAnalysisResponse
            on { listCodeAnalysisFindings(any<ListCodeAnalysisFindingsRequest>()) } doReturn listCodeAnalysisFindingsResponse
        }

        val mockConnection = mock<BearerSsoConnection>()
        whenever(mockConnection.getConnectionSettings()) doReturn mock<TokenConnectionSettings>()

        connectionManager = mock {
            on {
                activeConnectionForFeature(any())
            } doReturn authManagerRule.createConnection(ManagedSsoProfile("us-east-1", aString(), emptyList())) as AwsBearerTokenConnection
        }
        projectRule.project.replaceService(ToolkitConnectionManager::class.java, connectionManager, disposableRule.disposable)
    }

    @Test
    fun `Sono region is us-east-1`() {
        assertThat("us-east-1").isEqualTo(SONO_REGION)
    }

    @Test
    fun `generateCompletionsPaginator - bearer`() {
        val request = pythonRequest
        bearerClient.stub { client ->
            on { client.generateCompletions(any<GenerateCompletionsRequest>()) } doReturnConsecutively listOf(
                pythonResponseWithToken("first"),
                pythonResponseWithToken("second"),
                pythonResponseWithToken(""),
            )
        }

        val nextTokens = listOf("first", "second", "")
        val responses = sut.generateCompletionsPaginator(request)

        argumentCaptor<GenerateCompletionsRequest>().apply {
            responses.forEachIndexed { i, response ->
                assertThat(response.nextToken()).isEqualTo(nextTokens[i])
                response.completions().forEachIndexed { j, recommendation ->
                    assertThat(recommendation)
                        .usingRecursiveComparison()
                        .isEqualTo(response.completions()[j])
                }
            }
            verify(bearerClient, times(3)).generateCompletions(capture())
            verifyNoInteractions(sigv4Client)
            assertThat(this.firstValue.nextToken()).isEqualTo("")
            assertThat(this.secondValue.nextToken()).isEqualTo("first")
            assertThat(this.thirdValue.nextToken()).isEqualTo("second")
        }
    }

    @Test
    fun `createUploadUrl - bearer`() {
        val actual = sut.createUploadUrl(createUploadUrlRequest)

        argumentCaptor<CreateUploadUrlRequest>().apply {
            verify(bearerClient).createUploadUrl(capture())
            verifyNoInteractions(sigv4Client)
            assertThat(actual).isInstanceOf(CreateUploadUrlResponse::class.java)
            assertThat(actual).usingRecursiveComparison()
                .comparingOnlyFields("uploadUrl", "uploadId")
                .isEqualTo(createUploadUrlResponse)
        }
    }

    @Test
    fun `createCodeScan - sigv4`() {
        val actual = sut.createCodeScan(createCodeScanRequest, true)

        argumentCaptor<CreateCodeScanRequest>().apply {
            verify(sigv4Client).createCodeScan(capture())
            verifyNoInteractions(bearerClient)
            assertThat(firstValue).isSameAs(createCodeScanRequest)
            assertThat(actual).isSameAs(createCodeScanResponse)
        }
    }

    @Test
    fun `createCodeScan - bearer`() {
        val actual = sut.createCodeScan(createCodeScanRequest, false)

        argumentCaptor<StartCodeAnalysisRequest>().apply {
            verify(bearerClient).startCodeAnalysis(capture())
            verifyNoInteractions(sigv4Client)
            assertThat(actual).isInstanceOf(CreateCodeScanResponse::class.java)
            assertThat(actual).usingRecursiveComparison()
                .comparingOnlyFields("jobId", "status", "errorMessage")
                .isEqualTo(startCodeAnalysisResponse)
        }
    }

    @Test
    fun `getCodeScan - sigv4`() {
        val actual = sut.getCodeScan(getCodeScanRequest, true)

        argumentCaptor<GetCodeScanRequest>().apply {
            verify(sigv4Client).getCodeScan(capture())
            verifyNoInteractions(bearerClient)
            assertThat(firstValue).isSameAs(getCodeScanRequest)
            assertThat(actual).isSameAs(getCodeScanResponse)
        }
    }

    @Test
    fun `getCodeScan - bearer`() {
        val actual = sut.getCodeScan(getCodeScanRequest, false)

        argumentCaptor<GetCodeAnalysisRequest>().apply {
            verify(bearerClient).getCodeAnalysis(capture())
            verifyNoInteractions(sigv4Client)
            assertThat(actual).isInstanceOf(GetCodeScanResponse::class.java)
            assertThat(actual).usingRecursiveComparison()
                .comparingOnlyFields("status", "errorMessage")
                .isEqualTo(getCodeAnalysisResponse)
        }
    }

    @Test
    fun `listCodeScanFindings - sigv4`() {
        val actual = sut.listCodeScanFindings(listCodeScanFindingsRequest, true)

        argumentCaptor<ListCodeScanFindingsRequest>().apply {
            verify(sigv4Client).listCodeScanFindings(capture())
            verifyNoInteractions(bearerClient)
            assertThat(firstValue).isSameAs(listCodeScanFindingsRequest)
            assertThat(actual).isSameAs(listCodeScanFindingsResponse)
        }
    }

    @Test
    fun `listCodeScanFindings - bearer`() {
        val actual = sut.listCodeScanFindings(listCodeScanFindingsRequest, false)

        argumentCaptor<ListCodeAnalysisFindingsRequest>().apply {
            verify(bearerClient).listCodeAnalysisFindings(capture())
            verifyNoInteractions(sigv4Client)
            assertThat(actual).isInstanceOf(ListCodeScanFindingsResponse::class.java)
            assertThat(actual.codeScanFindings()).isEqualTo(listCodeAnalysisFindingsResponse.codeAnalysisFindings())
            assertThat(actual.nextToken()).isEqualTo(listCodeAnalysisFindingsResponse.nextToken())
        }
    }

    private companion object {
        val createCodeScanRequest = CreateCodeScanRequest.builder()
            .artifacts(mapOf(ArtifactType.SOURCE_CODE to "foo"))
            .clientToken("token")
            .programmingLanguage(
                ProgrammingLanguage.builder()
                    .languageName("python")
                    .build()
            )
            .build()

        val createUploadUrlRequest = CreateUploadUrlRequest.builder()
            .contentMd5("foo")
            .artifactType(software.amazon.awssdk.services.codewhispererruntime.model.ArtifactType.SOURCE_CODE)
            .build()

        val getCodeScanRequest = GetCodeScanRequest.builder()
            .jobId("jobid")
            .build()

        val listCodeScanFindingsRequest = ListCodeScanFindingsRequest.builder()
            .codeScanFindingsSchema(CodeScanFindingsSchema.CODESCAN_FINDINGS_1_0)
            .jobId("listCodeScanFindings - JobId")
            .nextToken("nextToken")
            .build()

        val createUploadUrlResponse: CreateUploadUrlResponse = CreateUploadUrlResponse.builder()
            .uploadUrl("url")
            .uploadId("id")
            .responseMetadata(metadata)
            .sdkHttpResponse(sdkHttpResponse)
            .build() as CreateUploadUrlResponse

        val startCodeAnalysisResponse = StartCodeAnalysisResponse.builder()
            .jobId("create-code-scan-user")
            .status(CodeAnalysisStatus.COMPLETED)
            .errorMessage("message")
            .responseMetadata(metadata)
            .sdkHttpResponse(sdkHttpResponse)
            .build() as StartCodeAnalysisResponse

        val getCodeAnalysisResponse = GetCodeAnalysisResponse.builder()
            .status(CodeAnalysisStatus.PENDING)
            .errorMessage("message")
            .responseMetadata(metadata)
            .sdkHttpResponse(sdkHttpResponse)
            .build() as GetCodeAnalysisResponse

        val listCodeAnalysisFindingsResponse = ListCodeAnalysisFindingsResponse.builder()
            .codeAnalysisFindings("findings")
            .nextToken("nextToken")
            .responseMetadata(metadata)
            .sdkHttpResponse(sdkHttpResponse)
            .build() as ListCodeAnalysisFindingsResponse

        private val generateCompletionsPaginatorResponse: GenerateCompletionsIterable = mock()

        private val createCodeScanUploadUrlResponse: CreateCodeScanUploadUrlResponse = mock()

        private val createCodeScanResponse: CreateCodeScanResponse = mock()

        private val getCodeScanResponse: GetCodeScanResponse = mock()

        private val listCodeScanFindingsResponse: ListCodeScanFindingsResponse = mock()
    }
}