From 76ed5537567e37a07c691cbb6c2eea9932fcb7c9 Mon Sep 17 00:00:00 2001
From: Pavel Zhukov <33721692+LeaveMyYard@users.noreply.github.com>
Date: Thu, 4 Apr 2024 13:41:07 +0300
Subject: Refactor k8s workloads streaming (#256)

* Refactor k8s workloads streaming

* Fix tests
---
 tests/conftest.py | 13 ++-----------
 1 file changed, 2 insertions(+), 11 deletions(-)

(limited to 'tests')

diff --git a/tests/conftest.py b/tests/conftest.py
index 61c389d..b1d8d22 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,6 +1,6 @@
 import random
 from datetime import datetime, timedelta
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import AsyncMock, patch
 
 import numpy as np
 import pytest
@@ -26,15 +26,6 @@ TEST_OBJECT = K8sObjectData(
 )
 
 
-class AsyncIter:
-    def __init__(self, items):
-        self.items = items
-
-    async def __aiter__(self):
-        for item in self.items:
-            yield item
-
-
 @pytest.fixture(autouse=True, scope="session")
 def mock_list_clusters():
     with patch(
@@ -48,7 +39,7 @@ def mock_list_clusters():
 def mock_list_scannable_objects():
     with patch(
         "robusta_krr.core.integrations.kubernetes.KubernetesLoader.list_scannable_objects",
-        new=MagicMock(return_value=AsyncIter([TEST_OBJECT])),
+        new=AsyncMock(return_value=[TEST_OBJECT]),
     ):
         yield
 
-- 
cgit v1.2.3