Make tests pytest and unittest compatible (#254)

Author: codesue

model_card_toolkit/core_test.py

@@ -16,6 +16,7 @@
 import os
 from unittest import mock
 
+from absl import flags
 from absl.testing import absltest
 from absl.testing import parameterized
 
@@ -289,3 +290,7 @@ def test_export_format_before_scaffold_assets(self):
 
 if __name__ == '__main__':
   absltest.main()
+else:
+  # Manually pass and parse flags to prevent UnparsedFlagAccessError when using
+  # pytest or unittest as a runner.
+  flags.FLAGS(['--test_tmpdir'])

model_card_toolkit/utils/testdata/tfxtest.py

@@ -16,6 +16,7 @@
 import os
 from typing import Any, Callable, List, Optional
 
+from absl import flags
 import apache_beam as beam
 from model_card_toolkit.utils.tfx_util import _TFX_METRICS_TYPE
 from model_card_toolkit.utils.tfx_util import _TFX_STATS_TYPE
@@ -174,3 +175,9 @@ def _write(dataset_name: str, features: List[str], split_name: str):
 
     if store:
       self._put_artifact(store, _TFX_STATS_TYPE, tfdv_path)
+
+
+if not __name__ == '__main__':
+  # Manually pass and parse flags to prevent UnparsedFlagAccessError when using
+  # pytest or unittest as a runner.
+  flags.FLAGS(['--test_tmpdir'])