From af153d8ccbf87fa890a6e107ab35ab872c18c72e Mon Sep 17 00:00:00 2001
From: Pileks <jure@amber-it.co>
Date: Fri, 22 Sep 2023 17:29:39 +0200
Subject: [PATCH] fix test code

---
 .../plugins/polywrap-logger-plugin/tests/conftest.py |  6 +++---
 .../tests/integration/test_log.py                    |  6 +++---
 .../polywrap-logger-plugin/tests/unit/test_log.py    | 12 ++++++------
 3 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/packages/plugins/polywrap-logger-plugin/tests/conftest.py b/packages/plugins/polywrap-logger-plugin/tests/conftest.py
index b020e7ff..753dbea6 100644
--- a/packages/plugins/polywrap-logger-plugin/tests/conftest.py
+++ b/packages/plugins/polywrap-logger-plugin/tests/conftest.py
@@ -1,12 +1,12 @@
 import logging
 import pytest
-from polywrap_logger_plugin import LoggerConfig, LoggerLogLevel
+from polywrap_logger_plugin import LoggerConfig, LogLevel
 from typing import Callable
 
 
 @pytest.fixture
-def valid_logger_config() -> Callable[[LoggerLogLevel], LoggerConfig]:
-    def _valid_logger_config(level: LoggerLogLevel) -> LoggerConfig:
+def valid_logger_config() -> Callable[[LogLevel], LoggerConfig]:
+    def _valid_logger_config(level: LogLevel) -> LoggerConfig:
         return LoggerConfig(logger=logging.getLogger("test_logger"), level=level)
     return _valid_logger_config
 
diff --git a/packages/plugins/polywrap-logger-plugin/tests/integration/test_log.py b/packages/plugins/polywrap-logger-plugin/tests/integration/test_log.py
index 7cfdda6e..98114735 100644
--- a/packages/plugins/polywrap-logger-plugin/tests/integration/test_log.py
+++ b/packages/plugins/polywrap-logger-plugin/tests/integration/test_log.py
@@ -2,7 +2,7 @@
 from polywrap_client import PolywrapClient
 from polywrap_core import Uri
 import pytest
-from polywrap_logger_plugin import LoggerLogLevel
+from polywrap_logger_plugin import LogLevel
 from _pytest.logging import LogCaptureFixture
 
 
@@ -10,7 +10,7 @@
 async def test_log_levels(
     client: PolywrapClient,
     caplog: LogCaptureFixture,
-    log_level: LoggerLogLevel,
+    log_level: LogLevel,
 ):
     caplog.set_level(logging.DEBUG)
     args_log = dict(message="Test message", level=log_level)
@@ -22,7 +22,7 @@ async def test_log_levels(
     )
     last_record = caplog.records[-1] if caplog.records else None
 
-    if log_level >= LoggerLogLevel.INFO:
+    if log_level >= LogLevel.INFO:
         assert last_record
         assert last_record.levelno == (log_level + 1) * 10
         assert last_record.message == "Test message"
diff --git a/packages/plugins/polywrap-logger-plugin/tests/unit/test_log.py b/packages/plugins/polywrap-logger-plugin/tests/unit/test_log.py
index a229825c..372ff4d6 100644
--- a/packages/plugins/polywrap-logger-plugin/tests/unit/test_log.py
+++ b/packages/plugins/polywrap-logger-plugin/tests/unit/test_log.py
@@ -1,18 +1,18 @@
 import logging
 from polywrap_core import InvokerClient
 import pytest
-from polywrap_logger_plugin import LoggerConfig, LoggerModule, ArgsLog, LoggerLogLevel
+from polywrap_logger_plugin import LoggerConfig, LoggerModule, ArgsLog, LogLevel
 from _pytest.logging import LogCaptureFixture
 from typing import Callable, cast
 
 
-@pytest.mark.parametrize("config_level", list(LoggerLogLevel))
-@pytest.mark.parametrize("log_level", list(LoggerLogLevel))
+@pytest.mark.parametrize("config_level", list(LogLevel))
+@pytest.mark.parametrize("log_level", list(LogLevel))
 async def test_log_levels(
     caplog: LogCaptureFixture,
-    valid_logger_config: Callable[[LoggerLogLevel], LoggerConfig],
-    config_level: LoggerLogLevel,
-    log_level: LoggerLogLevel,
+    valid_logger_config: Callable[[LogLevel], LoggerConfig],
+    config_level: LogLevel,
+    log_level: LogLevel,
 ):
     caplog.set_level(logging.DEBUG)
     logger_module = LoggerModule(valid_logger_config(config_level))