Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix Jira Metrics Behavior #5360

Closed
wants to merge 4 commits into from
Closed

Fix Jira Metrics Behavior #5360

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
55e5ecb
Me when i am in a metrics competition and my opponent is the metrics …
Galactus22625 Jan 27, 2025
5dfefd3
Merge branch 'main' into fix-metrics-then-cherry-pick
Galactus22625 Jan 27, 2025
7efbcae
me on cleanup island (cleaning up files and fixing tests)
Galactus22625 Jan 27, 2025
4f619dd
Merge branch 'main' into fix-metrics-then-cherry-pick
Galactus22625 Jan 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions ...jira-source/src/main/java/org/opensearch/dataprepper/plugins/source/jira/JiraService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,13 +60,16 @@ public class JiraService {

private final JiraSourceConfig jiraSourceConfig;
private final JiraRestClient jiraRestClient;
private final Counter searchResultsFoundCounter;
private final PluginMetrics jiraPluginMetrics = PluginMetrics.fromNames("jiraService", "aws");

private Counter searchResultsFoundCounter;
private PluginMetrics jiraPluginMetrics;

public JiraService(JiraSourceConfig jiraSourceConfig, JiraRestClient jiraRestClient) {
this.jiraSourceConfig = jiraSourceConfig;
this.jiraRestClient = jiraRestClient;
}

public void setPluginMetrics(PluginMetrics pluginMetrics){
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can pluginMetrics be added as part of JiraService constructor line 66 ?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tried that and it doesnt work because JiraService is a bean I think

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we remove the the bean annotation for JiraService ?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that would cause a lot of problems? not sure tho have to ask @san81

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@dinujoh we are using Spring dependency injection here. We are not instantiating JiraService ourself. Spring DI is instantiating it and for Spring DI, any constructor parameter should also be a spring bean otherwise, it won't inject as the constructor argument. If PluginMetrics is a spring bean then constructor argument will work but it is not.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

https://github.com/opensearch-project/data-prepper/blob/main/data-prepper-plugins/saas-source-plugins/jira-source/src/main/java/org/opensearch/dataprepper/plugins/source/jira/JiraService.java

Does this needs to use Spring dependency injection ? I was recommending to create the instance when needed in the client for example in JiraSource without using Spring DI

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why can't we include PluginMetrics in the context so that we can provide it via the application context? That would be a good long-term solution.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The idea is that, as we add additional SAAS sources, the common code in source-crawler will dynamically bind with one of the implementation in a specific SAAS source like Jira or Confluence etc. That is why we started using Spring DI in these source plugins.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@dlvenable Yah, I like the approach to include PluginMetrics in the application context. We will try that approach and update here.

this.jiraPluginMetrics = pluginMetrics;
this.searchResultsFoundCounter = jiraPluginMetrics.counter(SEARCH_RESULTS_FOUND);
}

Expand Down
9 changes: 8 additions & 1 deletion .../jira-source/src/main/java/org/opensearch/dataprepper/plugins/source/jira/JiraSource.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import org.opensearch.dataprepper.model.plugin.PluginFactory;
import org.opensearch.dataprepper.model.record.Record;
import org.opensearch.dataprepper.model.source.Source;
import org.opensearch.dataprepper.plugins.source.jira.rest.JiraRestClient;
import org.opensearch.dataprepper.plugins.source.jira.rest.auth.JiraAuthConfig;
import org.opensearch.dataprepper.plugins.source.jira.utils.JiraConfigHelper;
import org.opensearch.dataprepper.plugins.source.source_crawler.CrawlerApplicationContextMarker;
Expand Down Expand Up @@ -54,9 +55,15 @@ public JiraSource(final PluginMetrics pluginMetrics,
final PluginFactory pluginFactory,
final AcknowledgementSetManager acknowledgementSetManager,
Crawler crawler,
PluginExecutorServiceProvider executorServiceProvider) {
PluginExecutorServiceProvider executorServiceProvider,
final JiraService jiraService,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How is this injected here ? Can JiraService be initialized in the JiraSource constructor ? and pass pluginMetrics to the JiraService constructor ?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

tried passing in pluginMetrics to jiraService constructor but doesn't work because jira service is a bean i think. Jira Service can be initialized in jiraSource constructor for the same reason

final JiraRestClient jiraRestClient) {
super(PLUGIN_NAME, pluginMetrics, jiraSourceConfig, pluginFactory, acknowledgementSetManager, crawler, executorServiceProvider);
log.info("Creating Jira Source Plugin");
jiraService.setPluginMetrics(pluginMetrics);
crawler.setPluginMetrics(pluginMetrics);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same comment for crawler initialization.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same as jira service. Also crawler will be shared by multiple source plugins in the future

jiraRestClient.setPluginMetrics(pluginMetrics);

this.jiraSourceConfig = jiraSourceConfig;
this.jiraOauthConfig = jiraOauthConfig;
}
Expand Down
1 change: 1 addition & 0 deletions ...source/src/main/java/org/opensearch/dataprepper/plugins/source/jira/JiraSourceConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ public class JiraSourceConfig implements CrawlerSourceConfig {
@JsonProperty("acknowledgments")
private boolean acknowledgments = false;


public String getAccountUrl() {
return this.getHosts().get(0);
}
Expand Down
13 changes: 8 additions & 5 deletions ...rce/src/main/java/org/opensearch/dataprepper/plugins/source/jira/rest/JiraRestClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,17 +55,20 @@ public class JiraRestClient {
private static final String ISSUES_REQUESTED = "issuesRequested";
private final RestTemplate restTemplate;
private final JiraAuthConfig authConfig;
private final Timer ticketFetchLatencyTimer;
private final Timer searchCallLatencyTimer;
private final Timer projectFetchLatencyTimer;
private final Counter issuesRequestedCounter;
private final PluginMetrics jiraPluginMetrics = PluginMetrics.fromNames("jiraRestClient", "aws");
private Timer ticketFetchLatencyTimer;
private Timer searchCallLatencyTimer;
private Timer projectFetchLatencyTimer;
private Counter issuesRequestedCounter;
private PluginMetrics jiraPluginMetrics;
private int sleepTimeMultiplier = 1000;

public JiraRestClient(RestTemplate restTemplate, JiraAuthConfig authConfig) {
this.restTemplate = restTemplate;
this.authConfig = authConfig;
}

public void setPluginMetrics(PluginMetrics pluginMetrics){
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can pluginMetrics be passed in JiraRestClient constructor ?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No for the same reason as Jira Service

this.jiraPluginMetrics = pluginMetrics;
ticketFetchLatencyTimer = jiraPluginMetrics.timer(TICKET_FETCH_LATENCY_TIMER);
searchCallLatencyTimer = jiraPluginMetrics.timer(SEARCH_CALL_LATENCY_TIMER);
projectFetchLatencyTimer = jiraPluginMetrics.timer(PROJECTS_FETCH_LATENCY_TIMER);
Expand Down
13 changes: 13 additions & 0 deletions ...-source/src/test/java/org/opensearch/dataprepper/plugins/source/jira/JiraServiceTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,11 +13,13 @@
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
import io.micrometer.core.instrument.Counter;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.opensearch.dataprepper.metrics.PluginMetrics;
import org.opensearch.dataprepper.model.plugin.PluginConfigVariable;
import org.opensearch.dataprepper.plugins.source.jira.configuration.Oauth2Config;
import org.opensearch.dataprepper.plugins.source.jira.exception.BadRequestException;
Expand Down Expand Up @@ -76,6 +78,12 @@ public class JiraServiceTest {
@Mock
private JiraRestClient jiraRestClient;

@Mock
private PluginMetrics pluginMetrics;

@Mock
private Counter counter;


private static InputStream getResourceAsStream(String resourceName) {
InputStream inputStream = Thread.currentThread().getContextClassLoader().getResourceAsStream(resourceName);
Expand Down Expand Up @@ -210,6 +218,9 @@ public void testGetJiraEntities() throws JsonProcessingException {

Instant timestamp = Instant.ofEpochSecond(0);
Queue<ItemInfo> itemInfoQueue = new ConcurrentLinkedQueue<>();

when(pluginMetrics.counter(anyString())).thenReturn(counter);
jiraService.setPluginMetrics(pluginMetrics);
jiraService.getJiraEntities(jiraSourceConfig, timestamp, itemInfoQueue);
assertEquals(mockIssues.size(), itemInfoQueue.size());
}
Expand All @@ -236,6 +247,8 @@ public void buildIssueItemInfoMultipleFutureThreads() throws JsonProcessingExcep

Instant timestamp = Instant.ofEpochSecond(0);
Queue<ItemInfo> itemInfoQueue = new ConcurrentLinkedQueue<>();
when(pluginMetrics.counter(anyString())).thenReturn(counter);
jiraService.setPluginMetrics(pluginMetrics);
jiraService.getJiraEntities(jiraSourceConfig, timestamp, itemInfoQueue);
assertTrue(itemInfoQueue.size() >= 100);
}
Expand Down
15 changes: 11 additions & 4 deletions ...a-source/src/test/java/org/opensearch/dataprepper/plugins/source/jira/JiraSourceTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator;
import org.opensearch.dataprepper.plugins.source.jira.configuration.AuthenticationConfig;
import org.opensearch.dataprepper.plugins.source.jira.configuration.BasicConfig;
import org.opensearch.dataprepper.plugins.source.jira.rest.JiraRestClient;
import org.opensearch.dataprepper.plugins.source.jira.rest.auth.JiraAuthConfig;
import org.opensearch.dataprepper.plugins.source.source_crawler.base.Crawler;
import org.opensearch.dataprepper.plugins.source.source_crawler.base.PluginExecutorServiceProvider;
Expand Down Expand Up @@ -77,17 +78,23 @@ public class JiraSourceTest {
@Mock
BasicConfig basicConfig;

@Mock
private JiraService jiraService;

@Mock
private JiraRestClient jiraRestClient;

@Test
void initialization() {
when(executorServiceProvider.get()).thenReturn(executorService);
JiraSource source = new JiraSource(pluginMetrics, jiraSourceConfig, jiraOauthConfig, pluginFactory, acknowledgementSetManager, crawler, executorServiceProvider);
JiraSource source = new JiraSource(pluginMetrics, jiraSourceConfig, jiraOauthConfig, pluginFactory, acknowledgementSetManager, crawler, executorServiceProvider, jiraService, jiraRestClient);
assertNotNull(source);
}

@Test
void testStart() {
when(executorServiceProvider.get()).thenReturn(executorService);
JiraSource source = new JiraSource(pluginMetrics, jiraSourceConfig, jiraOauthConfig, pluginFactory, acknowledgementSetManager, crawler, executorServiceProvider);
JiraSource source = new JiraSource(pluginMetrics, jiraSourceConfig, jiraOauthConfig, pluginFactory, acknowledgementSetManager, crawler, executorServiceProvider, jiraService, jiraRestClient);
when(jiraSourceConfig.getAccountUrl()).thenReturn(ACCESSIBLE_RESOURCES);
when(jiraSourceConfig.getAuthType()).thenReturn(BASIC);
when(jiraSourceConfig.getAuthenticationConfig()).thenReturn(authenticationConfig);
Expand All @@ -103,7 +110,7 @@ void testStart() {
@Test
void testStop() {
when(executorServiceProvider.get()).thenReturn(executorService);
JiraSource source = new JiraSource(pluginMetrics, jiraSourceConfig, jiraOauthConfig, pluginFactory, acknowledgementSetManager, crawler, executorServiceProvider);
JiraSource source = new JiraSource(pluginMetrics, jiraSourceConfig, jiraOauthConfig, pluginFactory, acknowledgementSetManager, crawler, executorServiceProvider, jiraService, jiraRestClient);
when(jiraSourceConfig.getAccountUrl()).thenReturn(ACCESSIBLE_RESOURCES);
when(jiraSourceConfig.getAuthType()).thenReturn(BASIC);
when(jiraSourceConfig.getAuthenticationConfig()).thenReturn(authenticationConfig);
Expand All @@ -120,7 +127,7 @@ void testStop() {
@Test
void testStop_WhenNotStarted() {
when(executorServiceProvider.get()).thenReturn(executorService);
JiraSource source = new JiraSource(pluginMetrics, jiraSourceConfig, jiraOauthConfig, pluginFactory, acknowledgementSetManager, crawler, executorServiceProvider);
JiraSource source = new JiraSource(pluginMetrics, jiraSourceConfig, jiraOauthConfig, pluginFactory, acknowledgementSetManager, crawler, executorServiceProvider, jiraService, jiraRestClient);

source.stop();

Expand Down
31 changes: 31 additions & 0 deletions ...src/test/java/org/opensearch/dataprepper/plugins/source/jira/rest/JiraRestClientTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@
package org.opensearch.dataprepper.plugins.source.jira.rest;

import com.fasterxml.jackson.core.JsonProcessingException;
import io.micrometer.core.instrument.Counter;
import io.micrometer.core.instrument.Timer;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.junit.jupiter.params.ParameterizedTest;
Expand All @@ -19,6 +21,7 @@
import org.junit.jupiter.params.provider.ValueSource;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.opensearch.dataprepper.metrics.PluginMetrics;
import org.opensearch.dataprepper.plugins.source.jira.JiraServiceTest;
import org.opensearch.dataprepper.plugins.source.jira.JiraSourceConfig;
import org.opensearch.dataprepper.plugins.source.jira.exception.BadRequestException;
Expand All @@ -40,6 +43,7 @@
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
Expand All @@ -58,6 +62,15 @@ public class JiraRestClientTest {
@Mock
private JiraAuthConfig authConfig;

@Mock
private PluginMetrics pluginMetrics;

@Mock
private Timer timer;

@Mock
private Counter counter;

private static Stream<Arguments> provideHttpStatusCodesWithExceptionClass() {
return Stream.of(
Arguments.of(HttpStatus.FORBIDDEN, UnAuthorizedException.class),
Expand All @@ -75,6 +88,9 @@ public void testFetchingJiraIssue(String configFileName) {
JiraSourceConfig jiraSourceConfig = JiraServiceTest.createJiraConfigurationFromYaml(configFileName);
JiraAuthConfig authConfig = new JiraAuthFactory(jiraSourceConfig).getObject();
JiraRestClient jiraRestClient = new JiraRestClient(restTemplate, authConfig);
when(pluginMetrics.counter(anyString())).thenReturn(counter);
when(pluginMetrics.timer(anyString())).thenReturn(timer);
jiraRestClient.setPluginMetrics(pluginMetrics);
String ticketDetails = jiraRestClient.getIssue("key");
assertEquals(exampleTicketResponse, ticketDetails);
}
Expand All @@ -86,6 +102,9 @@ void testInvokeRestApiTokenExpired(HttpStatus statusCode, Class expectedExceptio
jiraRestClient.setSleepTimeMultiplier(1);
when(authConfig.getUrl()).thenReturn("https://example.com/rest/api/2/issue/key");
when(restTemplate.getForEntity(any(URI.class), any(Class.class))).thenThrow(new HttpClientErrorException(statusCode));
when(pluginMetrics.counter(anyString())).thenReturn(counter);
when(pluginMetrics.timer(anyString())).thenReturn(timer);
jiraRestClient.setPluginMetrics(pluginMetrics);
assertThrows(expectedExceptionType, () -> jiraRestClient.getIssue("key"));
}

Expand All @@ -94,6 +113,9 @@ void testInvokeRestApiTokenExpiredInterruptException() throws InterruptedExcepti
JiraRestClient jiraRestClient = new JiraRestClient(restTemplate, authConfig);
when(authConfig.getUrl()).thenReturn("https://example.com/rest/api/2/issue/key");
when(restTemplate.getForEntity(any(URI.class), any(Class.class))).thenThrow(new HttpClientErrorException(HttpStatus.TOO_MANY_REQUESTS));
when(pluginMetrics.counter(anyString())).thenReturn(counter);
when(pluginMetrics.timer(anyString())).thenReturn(timer);
jiraRestClient.setPluginMetrics(pluginMetrics);
jiraRestClient.setSleepTimeMultiplier(100000);

Thread testThread = new Thread(() -> {
Expand Down Expand Up @@ -121,6 +143,9 @@ public void testGetAllIssuesOauth2() throws JsonProcessingException {
SearchResults mockSearchResults = mock(SearchResults.class);
doReturn("http://mock-service.jira.com/").when(authConfig).getUrl();
doReturn(new ResponseEntity<>(mockSearchResults, HttpStatus.OK)).when(restTemplate).getForEntity(any(URI.class), any(Class.class));
when(pluginMetrics.counter(anyString())).thenReturn(counter);
when(pluginMetrics.timer(anyString())).thenReturn(timer);
jiraRestClient.setPluginMetrics(pluginMetrics);
SearchResults results = jiraRestClient.getAllIssues(jql, 0, jiraSourceConfig);
assertNotNull(results);
}
Expand All @@ -136,6 +161,9 @@ public void testGetAllIssuesBasic() throws JsonProcessingException {
SearchResults mockSearchResults = mock(SearchResults.class);
when(authConfig.getUrl()).thenReturn("https://example.com/");
doReturn(new ResponseEntity<>(mockSearchResults, HttpStatus.OK)).when(restTemplate).getForEntity(any(URI.class), any(Class.class));
when(pluginMetrics.counter(anyString())).thenReturn(counter);
when(pluginMetrics.timer(anyString())).thenReturn(timer);
jiraRestClient.setPluginMetrics(pluginMetrics);
SearchResults results = jiraRestClient.getAllIssues(jql, 0, jiraSourceConfig);
assertNotNull(results);
}
Expand All @@ -144,6 +172,9 @@ public void testGetAllIssuesBasic() throws JsonProcessingException {
public void testRestApiAddressValidation() throws JsonProcessingException {
when(authConfig.getUrl()).thenReturn("https://224.0.0.1/");
JiraRestClient jiraRestClient = new JiraRestClient(restTemplate, authConfig);
when(pluginMetrics.counter(anyString())).thenReturn(counter);
when(pluginMetrics.timer(anyString())).thenReturn(timer);
jiraRestClient.setPluginMetrics(pluginMetrics);
assertThrows(BadRequestException.class, () -> jiraRestClient.getIssue("TEST-1"));
}

Expand Down
10 changes: 6 additions & 4 deletions .../src/main/java/org/opensearch/dataprepper/plugins/source/source_crawler/base/Crawler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,14 +29,16 @@
@Named
public class Crawler {
private static final Logger log = LoggerFactory.getLogger(Crawler.class);
private final Timer crawlingTimer;
private final PluginMetrics pluginMetrics =
PluginMetrics.fromNames("sourceCrawler", "crawler");

private Timer crawlingTimer;
private PluginMetrics pluginMetrics;
private final CrawlerClient client;

public Crawler(CrawlerClient client) {
this.client = client;
}

public void setPluginMetrics(PluginMetrics pluginMetrics){
this.pluginMetrics = pluginMetrics;
this.crawlingTimer = pluginMetrics.timer("crawlingTime");
}

Expand Down
5 changes: 2 additions & 3 deletions ...rch/dataprepper/plugins/source/source_crawler/coordination/scheduler/WorkerScheduler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,6 @@ public class WorkerScheduler implements Runnable {
private final Counter acknowledgementSetSuccesses;
private final Counter acknowledgementSetFailures;
private final String sourcePluginName;
private final String SOURCE_PLUGIN_NAME = "sourcePluginName";


public WorkerScheduler(final String sourcePluginName,
Expand All @@ -58,8 +57,8 @@ public WorkerScheduler(final String sourcePluginName,

this.acknowledgementSetManager = acknowledgementSetManager;
this.pluginMetrics = pluginMetrics;
this.acknowledgementSetSuccesses = pluginMetrics.counterWithTags(ACKNOWLEDGEMENT_SET_SUCCESS_METRIC_NAME, SOURCE_PLUGIN_NAME, sourcePluginName);
this.acknowledgementSetFailures = pluginMetrics.counterWithTags(ACKNOWLEDGEMENT_SET_FAILURES_METRIC_NAME, SOURCE_PLUGIN_NAME, sourcePluginName);
this.acknowledgementSetSuccesses = pluginMetrics.counter(ACKNOWLEDGEMENT_SET_SUCCESS_METRIC_NAME);
this.acknowledgementSetFailures = pluginMetrics.counter(ACKNOWLEDGEMENT_SET_FAILURES_METRIC_NAME);
}

@Override
Expand Down
11 changes: 11 additions & 0 deletions .../test/java/org/opensearch/dataprepper/plugins/source/source_crawler/base/CrawlerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
package org.opensearch.dataprepper.plugins.source.source_crawler.base;

import io.micrometer.core.instrument.Timer;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.opensearch.dataprepper.metrics.PluginMetrics;
import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
import org.opensearch.dataprepper.model.buffer.Buffer;
import org.opensearch.dataprepper.model.event.Event;
Expand All @@ -27,6 +29,7 @@
import static org.junit.jupiter.api.Assertions.assertNotEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.reset;
import static org.mockito.Mockito.verify;
Expand All @@ -42,6 +45,12 @@ public class CrawlerTest {
private EnhancedSourceCoordinator coordinator;
@Mock
private Buffer<Record<Event>> buffer;
@Mock
private PluginMetrics pluginMetrics;

@Mock
private Timer timer;

@Mock
private CrawlerClient client;
@Mock
Expand All @@ -54,6 +63,8 @@ public class CrawlerTest {
@BeforeEach
public void setup() {
crawler = new Crawler(client);
when(pluginMetrics.timer(anyString())).thenReturn(timer);
crawler.setPluginMetrics(pluginMetrics);
when(leaderPartition.getProgressState()).thenReturn(Optional.of(new LeaderProgressState(lastPollTime)));
}

Expand Down
Loading