This is an automated email from the ASF dual-hosted git repository.
turbaszek pushed a commit to branch v1-10-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v1-10-test by this push:
new 547d03d Add airflow upgrade-check command (MVP) (#9467)
547d03d is described below
commit 547d03df2124e4345258d546e1c05707668119d6
Author: Kamil BreguĊa <[email protected]>
AuthorDate: Thu Sep 10 16:38:49 2020 +0200
Add airflow upgrade-check command (MVP) (#9467)
Co-authored-by: Tomek Urbaszek <[email protected]>
Co-authored-by: Kaxil Naik <[email protected]>
---
airflow/bin/cli.py | 31 +++++-
airflow/upgrade/__init__.py | 16 +++
airflow/upgrade/checker.py | 37 +++++++
airflow/upgrade/formatters.py | 114 +++++++++++++++++++++
airflow/upgrade/problem.py | 40 ++++++++
airflow/upgrade/rules/__init__.py | 16 +++
airflow/upgrade/rules/base_rule.py | 51 +++++++++
airflow/upgrade/rules/conn_type_is_not_nullable.py | 42 ++++++++
airflow/utils/cli.py | 23 +++++
docs/conf.py | 1 +
tests/upgrade/__init__.py | 16 +++
tests/upgrade/rules/__init__.py | 16 +++
tests/upgrade/rules/test_base_rule.py | 26 +++++
.../rules/test_conn_type_is_not_nullable.py | 40 ++++++++
tests/upgrade/test_formattes.py | 54 ++++++++++
tests/upgrade/test_problem.py | 35 +++++++
16 files changed, 553 insertions(+), 5 deletions(-)
diff --git a/airflow/bin/cli.py b/airflow/bin/cli.py
index fbc170d..45f1c5a 100644
--- a/airflow/bin/cli.py
+++ b/airflow/bin/cli.py
@@ -72,6 +72,8 @@ from airflow.models import (
)
from airflow.ti_deps.dep_context import (DepContext, SCHEDULER_QUEUED_DEPS)
from airflow.typing_compat import Protocol
+from airflow.upgrade.checker import check_upgrade
+from airflow.upgrade.formatters import (ConsoleFormatter, JSONFormatter)
from airflow.utils import cli as cli_utils, db
from airflow.utils.dot_renderer import render_dag
from airflow.utils.net import get_hostname
@@ -2273,6 +2275,19 @@ def info(args):
print(info)
+def upgrade_check(args):
+ if args.save:
+ filename = args.save
+ if not filename.lower().endswith(".json"):
+ print("Only JSON files are supported", file=sys.stderr)
+ formatter = JSONFormatter(args.save)
+ else:
+ formatter = ConsoleFormatter()
+ all_problems = check_upgrade(formatter)
+ if all_problems:
+ sys.exit(1)
+
+
class Arg(object):
def __init__(self, flags=None, help=None, action=None, default=None,
nargs=None,
type=None, choices=None, metavar=None):
@@ -2442,16 +2457,16 @@ class CLIFactory(object):
# show_dag
'save': Arg(
("-s", "--save"),
- "Saves the result to the indicated file.\n"
+ "Saves the result to the indicated file. The file format is
determined by the file extension.\n"
"\n"
- "The file format is determined by the file extension. For more
information about supported "
- "format, see: https://www.graphviz.org/doc/info/output.html\n";
+ "To see more information about supported format for show_dags
command, see: "
+ "https://www.graphviz.org/doc/info/output.html\n";
"\n"
"If you want to create a PNG file then you should execute the
following command:\n"
- "airflow dags show <DAG_ID> --save output.png\n"
+ "airflow show_dag <DAG_ID> --save output.png\n"
"\n"
"If you want to create a DOT file then you should execute the
following command:\n"
- "airflow dags show <DAG_ID> --save output.dot\n"
+ "airflow show_dag <DAG_ID> --save output.dot\n"
),
'imgcat': Arg(
("--imgcat", ),
@@ -3015,6 +3030,12 @@ class CLIFactory(object):
'func': info,
'args': ('anonymize', 'file_io', ),
},
+ {
+ 'name': 'upgrade_check',
+ 'help': 'Check if you can upgrade to the new version.',
+ 'func': upgrade_check,
+ 'args': ('save', ),
+ },
)
subparsers_dict = {sp['func'].__name__: sp for sp in subparsers}
dag_subparsers = (
diff --git a/airflow/upgrade/__init__.py b/airflow/upgrade/__init__.py
new file mode 100644
index 0000000..13a8339
--- /dev/null
+++ b/airflow/upgrade/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/airflow/upgrade/checker.py b/airflow/upgrade/checker.py
new file mode 100644
index 0000000..e8c6837
--- /dev/null
+++ b/airflow/upgrade/checker.py
@@ -0,0 +1,37 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import absolute_import
+from typing import List
+
+from airflow.upgrade.formatters import BaseFormatter
+from airflow.upgrade.problem import RuleStatus
+from airflow.upgrade.rules.base_rule import BaseRule, RULES
+
+ALL_RULES = [cls() for cls in RULES] # type: List[BaseRule]
+
+
+def check_upgrade(formatter):
+ # type: (BaseFormatter) -> List[RuleStatus]
+ formatter.start_checking(ALL_RULES)
+ all_rule_statuses = [] # List[RuleStatus]
+ for rule in ALL_RULES:
+ rule_status = RuleStatus.from_rule(rule)
+ all_rule_statuses.append(rule_status)
+ formatter.on_next_rule_status(rule_status)
+ formatter.end_checking(all_rule_statuses)
+ return all_rule_statuses
diff --git a/airflow/upgrade/formatters.py b/airflow/upgrade/formatters.py
new file mode 100644
index 0000000..be408dc
--- /dev/null
+++ b/airflow/upgrade/formatters.py
@@ -0,0 +1,114 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from abc import ABCMeta
+from typing import List
+import json
+
+from airflow.upgrade.problem import RuleStatus
+from airflow.upgrade.rules.base_rule import BaseRule
+from airflow.utils.cli import header, get_terminal_size
+
+
+class BaseFormatter(object):
+ __metaclass__ = ABCMeta
+
+ def start_checking(self, all_rules):
+ # type: (List[BaseRule]) -> None
+ pass
+
+ def end_checking(self, rule_statuses):
+ # type: (List[RuleStatus]) -> None
+
+ pass
+
+ def on_next_rule_status(self, rule_status):
+ # type: (RuleStatus) -> None
+ pass
+
+
+class ConsoleFormatter(BaseFormatter):
+ def start_checking(self, all_rules):
+ print()
+ header("STATUS", "=")
+ print()
+
+ def end_checking(self, rule_statuses):
+ if rule_statuses:
+ messages_count = sum(
+ len(rule_status.messages)
+ for rule_status in rule_statuses
+ )
+ if messages_count == 1:
+ print("Found {} problem.".format(messages_count))
+ else:
+ print("Found {} problems.".format(messages_count))
+ print()
+ header("RECOMMENDATIONS", "=")
+ print()
+
+ self.display_recommendations(rule_statuses)
+ else:
+ print("Not found any problems. World is beautiful. ")
+ print("You can safely update Airflow to the new version.")
+
+ @staticmethod
+ def display_recommendations(rule_statuses):
+
+ for rule_status in rule_statuses:
+ rule = rule_status.rule
+ print(rule.title)
+ print("-" * len(rule.title))
+ print(rule.description)
+ print("")
+ if rule_status.messages:
+ print("Problems:")
+ for message_no, message in enumerate(rule_status.messages, 1):
+ print('{:>3}. {}'.format(message_no, message))
+
+ def on_next_rule_status(self, rule_status):
+ status = "SUCCESS" if rule_status.is_success else "FAIL"
+ status_line_fmt = self.prepare_status_line_format()
+ print(status_line_fmt.format(rule_status.rule.title, status))
+
+ @staticmethod
+ def prepare_status_line_format():
+ _, terminal_width = get_terminal_size()
+
+ return "{:.<" + str(terminal_width - 10) + "}{:.>10}"
+
+
+class JSONFormatter(BaseFormatter):
+ def __init__(self, output_path):
+ self.filename = output_path
+
+ def start_checking(self, all_rules):
+ print("Start looking for problems.")
+
+ @staticmethod
+ def _info_from_rule_status(rule_status):
+ return {
+ "rule": type(rule_status.rule).__name__,
+ "title": rule_status.rule.title,
+ "messages": rule_status.messages,
+ }
+
+ def end_checking(self, rule_statuses):
+ formatted_results = [self._info_from_rule_status(rs) for rs in
rule_statuses]
+ with open(self.filename, "w+") as output_file:
+ json.dump(formatted_results, output_file, indent=2)
+ print("Saved result to: {}".format(self.filename))
diff --git a/airflow/upgrade/problem.py b/airflow/upgrade/problem.py
new file mode 100644
index 0000000..d5960e6
--- /dev/null
+++ b/airflow/upgrade/problem.py
@@ -0,0 +1,40 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import absolute_import
+from typing import NamedTuple, List
+
+from airflow.upgrade.rules.base_rule import BaseRule
+
+
+class RuleStatus(NamedTuple(
+ 'RuleStatus',
+ [
+ ('rule', BaseRule),
+ ('messages', List[str])
+ ]
+)):
+
+ @property
+ def is_success(self):
+ return bool(self.messages)
+
+ @classmethod
+ def from_rule(cls, rule):
+ # type: (BaseRule) -> RuleStatus
+ messages = rule.check()
+ return cls(rule=rule, messages=list(messages))
diff --git a/airflow/upgrade/rules/__init__.py
b/airflow/upgrade/rules/__init__.py
new file mode 100644
index 0000000..13a8339
--- /dev/null
+++ b/airflow/upgrade/rules/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/airflow/upgrade/rules/base_rule.py
b/airflow/upgrade/rules/base_rule.py
new file mode 100644
index 0000000..75ebe2f
--- /dev/null
+++ b/airflow/upgrade/rules/base_rule.py
@@ -0,0 +1,51 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from abc import ABCMeta, abstractmethod
+
+from six import add_metaclass
+
+RULES = []
+
+
+class BaseRuleMeta(ABCMeta):
+ def __new__(cls, clsname, bases, attrs):
+ clazz = super(BaseRuleMeta, cls).__new__(cls, clsname, bases, attrs)
+ if clsname != "BaseRule":
+ RULES.append(clazz)
+ return clazz
+
+
+@add_metaclass(BaseRuleMeta)
+class BaseRule(object):
+
+ @property
+ @abstractmethod
+ def title(self):
+ # type: () -> str
+ """Short one-line summary"""
+ pass
+
+ @property
+ @abstractmethod
+ def description(self):
+ # type: () -> str
+ """A long description explaining the problem in detail. This can be an
entry from UPDATING.md file."""
+ pass
+
+ def check(self):
+ pass
diff --git a/airflow/upgrade/rules/conn_type_is_not_nullable.py
b/airflow/upgrade/rules/conn_type_is_not_nullable.py
new file mode 100644
index 0000000..2ddd7c2
--- /dev/null
+++ b/airflow/upgrade/rules/conn_type_is_not_nullable.py
@@ -0,0 +1,42 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import absolute_import
+
+from airflow.models import Connection
+from airflow.upgrade.rules.base_rule import BaseRule
+from airflow.utils.db import provide_session
+
+
+class ConnTypeIsNotNullableRule(BaseRule):
+
+ title = "Connection.conn_type is not nullable"
+
+ description = """\
+The `conn_type` column in the `connection` table must contain content.
Previously, this rule was \
+enforced by application logic, but was not enforced by the database schema.
+
+If you made any modifications to the table directly, make sure you don't have
null in the conn_type column.\
+"""
+
+ @provide_session
+ def check(self, session=None):
+ invalid_connections =
session.query(Connection).filter(Connection.conn_type.is_(None))
+ return (
+ 'Connection<id={}", conn_id={}> have empty conn_type
field.'.format(conn.id, conn.conn_id)
+ for conn in invalid_connections
+ )
diff --git a/airflow/utils/cli.py b/airflow/utils/cli.py
index 334c58f..9018407 100644
--- a/airflow/utils/cli.py
+++ b/airflow/utils/cli.py
@@ -26,9 +26,13 @@ import functools
import getpass
import json
import socket
+import struct
import sys
from argparse import Namespace
from datetime import datetime
+from fcntl import ioctl
+from os import environ
+from termios import TIOCGWINSZ
from airflow.models import Log
from airflow.utils import cli_action_loggers
@@ -137,3 +141,22 @@ def should_use_colors(args):
if args.color == ColorMode.OFF:
return False
return is_terminal_support_colors()
+
+
+def get_terminal_size(fallback=(80, 20)):
+ """Return a tuple of (terminal height, terminal width)."""
+ try:
+ return struct.unpack('hhhh', ioctl(sys.__stdout__, TIOCGWINSZ, '\000'
* 8))[0:2]
+ except IOError:
+ # when the output stream or init descriptor is not a tty, such
+ # as when when stdout is piped to another program
+ pass
+ try:
+ return int(environ.get('LINES')), int(environ.get('COLUMNS'))
+ except TypeError:
+ return fallback
+
+
+def header(text, fillchar):
+ rows, columns = get_terminal_size()
+ print(" {} ".format(text).center(columns, fillchar))
diff --git a/docs/conf.py b/docs/conf.py
index 101d050..d70dd69 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -216,6 +216,7 @@ exclude_patterns = [
'_api/airflow/typing_compat',
'_api/airflow/kubernetes',
'_api/airflow/ti_deps',
+ '_api/airflow/upgrade',
'_api/airflow/utils',
'_api/airflow/version',
'_api/airflow/www',
diff --git a/tests/upgrade/__init__.py b/tests/upgrade/__init__.py
new file mode 100644
index 0000000..13a8339
--- /dev/null
+++ b/tests/upgrade/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/tests/upgrade/rules/__init__.py b/tests/upgrade/rules/__init__.py
new file mode 100644
index 0000000..13a8339
--- /dev/null
+++ b/tests/upgrade/rules/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/tests/upgrade/rules/test_base_rule.py
b/tests/upgrade/rules/test_base_rule.py
new file mode 100644
index 0000000..dbf47e3
--- /dev/null
+++ b/tests/upgrade/rules/test_base_rule.py
@@ -0,0 +1,26 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from airflow.upgrade.rules.base_rule import BaseRule, RULES
+
+
+class TestBaseRule:
+ def test_if_custom_rule_is_registered(self):
+ class CustomRule(BaseRule):
+ pass
+
+ assert CustomRule in list(RULES)
diff --git a/tests/upgrade/rules/test_conn_type_is_not_nullable.py
b/tests/upgrade/rules/test_conn_type_is_not_nullable.py
new file mode 100644
index 0000000..53a14a5
--- /dev/null
+++ b/tests/upgrade/rules/test_conn_type_is_not_nullable.py
@@ -0,0 +1,40 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from airflow.models import Connection
+from airflow.upgrade.rules.conn_type_is_not_nullable import
ConnTypeIsNotNullableRule
+from airflow.utils.db import create_session
+from tests.test_utils.db import clear_db_connections
+
+
+class TestConnTypeIsNotNullableRule:
+ def tearDown(self):
+ clear_db_connections()
+
+ def test_check(self):
+ rule = ConnTypeIsNotNullableRule()
+
+ assert isinstance(rule.description, str)
+ assert isinstance(rule.title, str)
+
+ with create_session() as session:
+ conn = Connection(conn_id="TestConnTypeIsNotNullableRule")
+ session.merge(conn)
+
+ msgs = rule.check(session=session)
+ assert [m for m in msgs if "TestConnTypeIsNotNullableRule" in m], \
+ "TestConnTypeIsNotNullableRule not in warning messages"
diff --git a/tests/upgrade/test_formattes.py b/tests/upgrade/test_formattes.py
new file mode 100644
index 0000000..0fc8f13
--- /dev/null
+++ b/tests/upgrade/test_formattes.py
@@ -0,0 +1,54 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import json
+from tempfile import NamedTemporaryFile
+from tests.compat import mock
+
+import pytest
+
+from airflow.bin import cli
+from airflow.upgrade.rules.base_rule import BaseRule
+
+MESSAGES = ["msg1", "msg2"]
+
+
+class MockRule(BaseRule):
+ title = "title"
+ description = "description"
+
+ def check(self):
+ return MESSAGES
+
+
+class TestJSONFormatter:
+ @mock.patch("airflow.upgrade.checker.ALL_RULES", [MockRule()])
+ def test_output(self):
+ expected = [
+ {
+ "rule": MockRule.__name__,
+ "title": MockRule.title,
+ "messages": MESSAGES,
+ }
+ ]
+ parser = cli.CLIFactory.get_parser()
+ with NamedTemporaryFile("w+") as temp:
+ with pytest.raises(SystemExit):
+ cli.upgrade_check(parser.parse_args(['upgrade_check', '-s',
temp.name]))
+ content = temp.read()
+
+ assert json.loads(content) == expected
diff --git a/tests/upgrade/test_problem.py b/tests/upgrade/test_problem.py
new file mode 100644
index 0000000..af2c5b9
--- /dev/null
+++ b/tests/upgrade/test_problem.py
@@ -0,0 +1,35 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from tests.compat import mock
+
+from airflow.upgrade.problem import RuleStatus
+
+
+class TestRuleStatus:
+ def test_is_success(self):
+ assert RuleStatus(rule=mock.MagicMock(), messages=[]).is_success is
False
+ assert RuleStatus(rule=mock.MagicMock(), messages=["aaa"]).is_success
is True
+
+ def test_rule_status_from_rule(self):
+ msgs = ["An interesting problem to solve"]
+ rule = mock.MagicMock()
+ rule.check.return_value = msgs
+
+ result = RuleStatus.from_rule(rule)
+ rule.check.assert_called_once_with()
+ assert result == RuleStatus(rule, msgs)
