1 files changed, 34 insertions, 28 deletions
diff --git a/spec/frontend/analytics/instance_statistics/apollo_mock_data.js b/spec/frontend/analytics/instance_statistics/apollo_mock_data.js
index 2e4eaf3fc96..98eabd577ee 100644
--- a/spec/frontend/analytics/instance_statistics/apollo_mock_data.js
+++ b/spec/frontend/analytics/instance_statistics/apollo_mock_data.js
@@ -1,30 +1,36 @@
-const defaultPageInfo = { hasPreviousPage: false, startCursor: null, endCursor: null };
+const defaultPageInfo = {
+  hasNextPage: false,
+  hasPreviousPage: false,
+  startCursor: null,
+  endCursor: null,
+};
 
-export function getApolloResponse(options = {}) {
-  const {
-    pipelinesTotal = [],
-    pipelinesSucceeded = [],
-    pipelinesFailed = [],
-    pipelinesCanceled = [],
-    pipelinesSkipped = [],
-    hasNextPage = false,
-  } = options;
-  return {
-    data: {
-      pipelinesTotal: { pageInfo: { ...defaultPageInfo, hasNextPage }, nodes: pipelinesTotal },
-      pipelinesSucceeded: {
-        pageInfo: { ...defaultPageInfo, hasNextPage },
-        nodes: pipelinesSucceeded,
-      },
-      pipelinesFailed: { pageInfo: { ...defaultPageInfo, hasNextPage }, nodes: pipelinesFailed },
-      pipelinesCanceled: {
-        pageInfo: { ...defaultPageInfo, hasNextPage },
-        nodes: pipelinesCanceled,
-      },
-      pipelinesSkipped: {
-        pageInfo: { ...defaultPageInfo, hasNextPage },
-        nodes: pipelinesSkipped,
-      },
+export const mockApolloResponse = ({ hasNextPage = false, key, data }) => ({
+  data: {
+    [key]: {
+      pageInfo: { ...defaultPageInfo, hasNextPage },
+      nodes: data,
     },
-  };
-}
+  },
+});
+
+export const mockQueryResponse = ({ key, data = [], loading = false, additionalData = [] }) => {
+  const hasNextPage = Boolean(additionalData.length);
+  const response = mockApolloResponse({ hasNextPage, key, data });
+  if (loading) {
+    return jest.fn().mockReturnValue(new Promise(() => {}));
+  }
+  if (hasNextPage) {
+    return jest
+      .fn()
+      .mockResolvedValueOnce(response)
+      .mockResolvedValueOnce(
+        mockApolloResponse({
+          hasNextPage: false,
+          key,
+          data: additionalData,
+        }),
+      );
+  }
+  return jest.fn().mockResolvedValue(response);
+};