coverage

Author: chakru-r

metadata-ingestion/tests/unit/cli/dataset/test_dataset_cmd.py

@@ -27,6 +27,31 @@ def test_yaml_file():
         temp_file.unlink()
 
 
+@pytest.fixture
+def invalid_value_yaml_file():
+    """Creates a temporary  yaml file - correctly formatted but bad datatype for testing."""
+    invalid_content = """
+## This file is intentionally malformed
+- id: user.badformat
+  platform: hive
+  schema:
+    fields:
+    - id: ip
+      type: bad_type
+      description: The IP address
+    """
+
+    # Create a temporary file
+    temp_file = TEST_RESOURCES_DIR / "invalid_dataset.yaml.tmp"
+    with open(temp_file, "w") as f:
+        f.write(invalid_content)
+
+    yield temp_file
+
+    # Clean up
+    if temp_file.exists():
+        temp_file.unlink()
+
 @pytest.fixture
 def malformed_yaml_file():
     """Creates a temporary malformed yaml file for testing."""
@@ -219,30 +244,29 @@ def test_multiple_datasets_in_file(self, mock_dataset, test_yaml_file):
             mock_dataset2.to_yaml.assert_called_once()
 
     @patch("datahub.cli.specific.dataset_cli.get_default_graph")
-    def test_dry_run_sync(self, mock_get_default_graph):
+    def test_dry_run_sync(self, mock_get_default_graph, test_yaml_file):
         mock_graph = MagicMock()
+        mock_graph.exists.return_value = True
         mock_get_default_graph.return_value.__enter__.return_value = mock_graph
 
-        test_file = TEST_RESOURCES_DIR / "dataset_dry_run.yaml"
-
         runner = CliRunner()
         result = runner.invoke(
-            dataset, ["sync", "--dry-run", "--to-datahub", "-f", str(test_file)]
+            dataset, ["sync", "--dry-run", "--to-datahub", "-f", str(test_yaml_file)]
         )
 
         # Verify
         assert result.exit_code == 0
         assert not mock_get_default_graph.emit.called
 
     @patch("datahub.cli.specific.dataset_cli.get_default_graph")
-    def test_dry_run_sync_fail(self, mock_get_default_graph):
+    def test_dry_run_sync_fail_bad_type(self, mock_get_default_graph, invalid_value_yaml_file):
         mock_graph = MagicMock()
+        mock_graph.exists.return_value = True
         mock_get_default_graph.return_value.__enter__.return_value = mock_graph
 
-        test_file = TEST_RESOURCES_DIR / "dataset_dry_run_bad_type.yaml"
         runner = CliRunner()
         result = runner.invoke(
-            dataset, ["sync", "--dry-run", "--to-datahub", "-f", str(test_file)]
+            dataset, ["sync", "--dry-run", "--to-datahub", "-f", str(invalid_value_yaml_file)]
         )
 
         # Verify
@@ -251,29 +275,72 @@ def test_dry_run_sync_fail(self, mock_get_default_graph):
         assert "Type bad_type is not a valid primitive type" in result.output
 
     @patch("datahub.cli.specific.dataset_cli.get_default_graph")
-    def test_run_sync(self, mock_get_default_graph):
+    def test_dry_run_sync_fail_missing_ref(self, mock_get_default_graph, test_yaml_file):
         mock_graph = MagicMock()
+        mock_graph.exists.return_value = False
         mock_get_default_graph.return_value.__enter__.return_value = mock_graph
 
-        test_file = TEST_RESOURCES_DIR / "dataset_dry_run.yaml"
+        runner = CliRunner()
+        result = runner.invoke(
+            dataset, ["sync", "--dry-run", "--to-datahub", "-f", str(test_yaml_file)]
+        )
+
+        # Verify
+        assert result.exit_code != 0
+        assert not mock_get_default_graph.emit.called
+        assert "missing entity reference" in result.output
+
+    @patch("datahub.cli.specific.dataset_cli.get_default_graph")
+    def test_run_sync(self, mock_get_default_graph, test_yaml_file):
+        mock_graph = MagicMock()
+        mock_graph.exists.return_value = True
+        mock_get_default_graph.return_value.__enter__.return_value = mock_graph
 
         runner = CliRunner()
-        result = runner.invoke(dataset, ["sync", "--to-datahub", "-f", str(test_file)])
+        result = runner.invoke(dataset, ["sync", "--to-datahub", "-f", str(test_yaml_file)])
 
         # Verify
         assert result.exit_code == 0
         assert mock_graph.emit.called
 
     @patch("datahub.cli.specific.dataset_cli.get_default_graph")
-    def test_run_sync_fail(self, mock_get_default_graph):
+    def test_run_sync_fail(self, mock_get_default_graph, invalid_value_yaml_file):
+        mock_graph = MagicMock()
+        mock_graph.exists.return_value = True
+        mock_get_default_graph.return_value.__enter__.return_value = mock_graph
+
+        runner = CliRunner()
+        result = runner.invoke(dataset, ["sync", "--to-datahub", "-f", str(invalid_value_yaml_file)])
+
+        # Verify
+        assert result.exit_code != 0
+        assert not mock_get_default_graph.emit.called
+        assert "is not a valid primitive type" in result.output
+
+    @patch("datahub.cli.specific.dataset_cli.get_default_graph")
+    def test_run_upsert_fail(self, mock_get_default_graph, invalid_value_yaml_file):
         mock_graph = MagicMock()
+        mock_graph.exists.return_value = True
         mock_get_default_graph.return_value.__enter__.return_value = mock_graph
 
-        test_file = TEST_RESOURCES_DIR / "dataset_dry_run_bad_type.yaml"
         runner = CliRunner()
-        result = runner.invoke(dataset, ["sync", "--to-datahub", "-f", str(test_file)])
+        result = runner.invoke(dataset, ["upsert", "-f", str(invalid_value_yaml_file)])
 
         # Verify
         assert result.exit_code != 0
         assert not mock_get_default_graph.emit.called
         assert "is not a valid primitive type" in result.output
+
+    @patch("datahub.cli.specific.dataset_cli.get_default_graph")
+    def test_sync_from_datahub_fail(self, mock_get_default_graph, test_yaml_file):
+        mock_graph = MagicMock()
+        mock_graph.exists.return_value = False
+        mock_get_default_graph.return_value.__enter__.return_value = mock_graph
+
+        runner = CliRunner()
+        result = runner.invoke(dataset, ["sync", "--dry-run", "--from-datahub", "-f", str(test_yaml_file)])
+
+        # Verify
+        assert result.exit_code != 0
+        assert "does not exist" in result.output
+