From 2d4a1841f65d96120abf24893ad72281afc3e21e Mon Sep 17 00:00:00 2001
From: Aarni Koskela <akx@iki.fi>
Date: Wed, 3 Nov 2021 19:33:57 +0200
Subject: [PATCH] Work around valohai-yaml compat bug

---
 tests/test_pipeline.py | 12 +++++-------
 tests/test_yaml.py     | 12 +++++-------
 tests/utils.py         | 12 ++++++++++++
 3 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py
index 5c898a1..4eef248 100644
--- a/tests/test_pipeline.py
+++ b/tests/test_pipeline.py
@@ -4,16 +4,15 @@
 import pytest
 from valohai_yaml import parse
 
-from tests.utils import read_yaml_test_data
+from tests.utils import compare_yaml, read_yaml_test_data
 from valohai.internals.pipeline import get_pipeline_from_source
-from valohai.yaml import config_to_yaml
 
 
 @pytest.mark.parametrize(
-    "original_yaml, source_python, expected_yaml",
+    "original_yaml, source_python, expected_yaml_filename",
     read_yaml_test_data("tests/test_pipeline_yaml"),
 )
-def test_yaml_update_from_source(tmpdir, original_yaml, source_python, expected_yaml):
+def test_pipeline_yaml_update_from_source(tmpdir, original_yaml, source_python, expected_yaml_filename):
     yaml_path = os.path.join(tmpdir, "valohai.yaml")
     source_path = os.path.join(tmpdir, "test.py")
 
@@ -33,6 +32,5 @@ def test_yaml_update_from_source(tmpdir, original_yaml, source_python, expected_
     new_config = old_config.merge_with(new_config)
 
     # Check against expected result
-    with open(expected_yaml) as expected_yaml:
-        new_yaml = config_to_yaml(new_config)
-        assert new_yaml == expected_yaml.read()
+    with open(expected_yaml_filename) as fp:
+        compare_yaml(new_config, fp.read())
diff --git a/tests/test_yaml.py b/tests/test_yaml.py
index a7242be..6029ddb 100644
--- a/tests/test_yaml.py
+++ b/tests/test_yaml.py
@@ -4,17 +4,16 @@
 import pytest
 from valohai_yaml import parse
 
-from tests.utils import read_yaml_test_data
+from tests.utils import compare_yaml, read_yaml_test_data
 from valohai.internals.merge import python_to_yaml_merge_strategy
 from valohai.internals.yaml import generate_step, parse_config_from_source
-from valohai.yaml import config_to_yaml
 
 
 @pytest.mark.parametrize(
-    "original_yaml, source_python, expected_yaml",
+    "original_yaml, source_python, expected_yaml_filename",
     read_yaml_test_data("tests/test_yaml"),
 )
-def test_yaml_update_from_source(tmpdir, original_yaml, source_python, expected_yaml):
+def test_yaml_update_from_source(tmpdir, original_yaml, source_python, expected_yaml_filename):
     yaml_path = os.path.join(tmpdir, "valohai.yaml")
     filename, file_extension = os.path.splitext(source_python)
     source_path = os.path.join(tmpdir, f"test{file_extension}")
@@ -38,9 +37,8 @@ def test_yaml_update_from_source(tmpdir, original_yaml, source_python, expected_
         new_config = old_config.merge_with(new_config, python_to_yaml_merge_strategy)
 
     # Check against expected result
-    with open(expected_yaml) as expected_yaml:
-        new_yaml = config_to_yaml(new_config)
-        assert new_yaml == expected_yaml.read()
+    with open(expected_yaml_filename) as fp:
+        compare_yaml(new_config, fp.read())
 
 
 def test_posix_path_separator(monkeypatch):
diff --git a/tests/utils.py b/tests/utils.py
index c5c8f3f..ff305c2 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,6 +1,11 @@
 import glob
 import json
 import os
+import re
+
+from valohai_yaml.objs import Config
+
+from valohai.yaml import config_to_yaml
 
 
 def read_source_files(path_without_ext):
@@ -56,3 +61,10 @@ def read_yaml_test_data(root_path):
             )
         )
     return test_data
+
+
+def compare_yaml(config: Config, fixture_yaml: str) -> None:
+    new_yaml = config_to_yaml(config)
+    # TODO: remove this when https://github.com/valohai/valohai-yaml/pull/59 lands
+    new_yaml = re.sub(r"\s+actions: \[]", "", new_yaml)
+    assert new_yaml == fixture_yaml