From 2491690c76c940169c1f636df04a4041ad59b80f Mon Sep 17 00:00:00 2001
From: Andre Anjos <andre.dos.anjos@gmail.com>
Date: Wed, 2 Aug 2023 21:01:31 +0200
Subject: [PATCH] [tests.conftest] Prepare helpers for multi-class database
support
---
tests/conftest.py | 12 ++++++++----
1 file changed, 8 insertions(+), 4 deletions(-)
diff --git a/tests/conftest.py b/tests/conftest.py
index 6e992cdd..82eef644 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -162,10 +162,14 @@ class DatabaseCheckers:
assert len(split[k]) == lengths[k]
for s in split[k]:
- assert any(
- [s[0].startswith(k) for k in prefixes]
- ), f"Sample with name {s[0]} does not start with any of the prefixes in {prefixes}"
- assert s[1] in possible_labels
+ assert any([s[0].startswith(k) for k in prefixes]), (
+ f"Sample with name {s[0]} does not start with any of the "
+ f"prefixes in {prefixes}"
+ )
+ if isinstance(s[1], list):
+ assert all([k in possible_labels for k in s[1]])
+ else:
+ assert s[1] in possible_labels
@staticmethod
def check_loaded_batch(
--
GitLab