|
37 | 37 | from unittest.mock import MagicMock |
38 | 38 | from kubernetes import client |
39 | 39 | import yaml |
| 40 | +import pytest |
40 | 41 | import filecmp |
41 | 42 | import os |
42 | 43 |
|
@@ -770,6 +771,136 @@ def custom_side_effect(group, version, namespace, plural, **kwargs): |
770 | 771 | assert result.dashboard == rc_dashboard |
771 | 772 |
|
772 | 773 |
|
| 774 | +def test_throw_for_no_raycluster_crd_errors(mocker): |
| 775 | + """Test RayCluster CRD error handling""" |
| 776 | + from kubernetes.client.rest import ApiException |
| 777 | + |
| 778 | + mocker.patch("kubernetes.config.load_kube_config", return_value="ignore") |
| 779 | + |
| 780 | + # Test 404 error - CRD not found |
| 781 | + mock_api_404 = MagicMock() |
| 782 | + mock_api_404.list_namespaced_custom_object.side_effect = ApiException(status=404) |
| 783 | + mocker.patch("kubernetes.client.CustomObjectsApi", return_value=mock_api_404) |
| 784 | + |
| 785 | + cluster = create_cluster(mocker) |
| 786 | + with pytest.raises( |
| 787 | + RuntimeError, match="RayCluster CustomResourceDefinition unavailable" |
| 788 | + ): |
| 789 | + cluster._throw_for_no_raycluster() |
| 790 | + |
| 791 | + # Test other API error |
| 792 | + mock_api_500 = MagicMock() |
| 793 | + mock_api_500.list_namespaced_custom_object.side_effect = ApiException(status=500) |
| 794 | + mocker.patch("kubernetes.client.CustomObjectsApi", return_value=mock_api_500) |
| 795 | + |
| 796 | + cluster2 = create_cluster(mocker) |
| 797 | + with pytest.raises( |
| 798 | + RuntimeError, match="Failed to get RayCluster CustomResourceDefinition" |
| 799 | + ): |
| 800 | + cluster2._throw_for_no_raycluster() |
| 801 | + |
| 802 | + |
| 803 | +def test_cluster_apply_attribute_error_handling(mocker): |
| 804 | + """Test AttributeError handling when DynamicClient fails""" |
| 805 | + mocker.patch("kubernetes.config.load_kube_config", return_value="ignore") |
| 806 | + mocker.patch("codeflare_sdk.ray.cluster.cluster.Cluster._throw_for_no_raycluster") |
| 807 | + mocker.patch( |
| 808 | + "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object", |
| 809 | + return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"), |
| 810 | + ) |
| 811 | + |
| 812 | + # Mock get_dynamic_client to raise AttributeError |
| 813 | + def raise_attribute_error(): |
| 814 | + raise AttributeError("DynamicClient initialization failed") |
| 815 | + |
| 816 | + mocker.patch( |
| 817 | + "codeflare_sdk.ray.cluster.cluster.Cluster.get_dynamic_client", |
| 818 | + side_effect=raise_attribute_error, |
| 819 | + ) |
| 820 | + |
| 821 | + cluster = create_cluster(mocker) |
| 822 | + |
| 823 | + with pytest.raises(RuntimeError, match="Failed to initialize DynamicClient"): |
| 824 | + cluster.apply() |
| 825 | + |
| 826 | + |
| 827 | +def test_cluster_namespace_handling(mocker, capsys): |
| 828 | + """Test namespace validation in create_resource""" |
| 829 | + mocker.patch("kubernetes.config.load_kube_config", return_value="ignore") |
| 830 | + mocker.patch( |
| 831 | + "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object", |
| 832 | + return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"), |
| 833 | + ) |
| 834 | + |
| 835 | + # Test with None namespace that gets set |
| 836 | + mocker.patch( |
| 837 | + "codeflare_sdk.ray.cluster.cluster.get_current_namespace", return_value=None |
| 838 | + ) |
| 839 | + |
| 840 | + config = ClusterConfiguration( |
| 841 | + name="test-cluster-ns", |
| 842 | + namespace=None, # Will trigger namespace check |
| 843 | + num_workers=1, |
| 844 | + worker_cpu_requests=1, |
| 845 | + worker_cpu_limits=1, |
| 846 | + worker_memory_requests=2, |
| 847 | + worker_memory_limits=2, |
| 848 | + ) |
| 849 | + |
| 850 | + cluster = Cluster(config) |
| 851 | + captured = capsys.readouterr() |
| 852 | + # Verify the warning message was printed |
| 853 | + assert "Please specify with namespace=<your_current_namespace>" in captured.out |
| 854 | + assert cluster.config.namespace is None |
| 855 | + |
| 856 | + |
| 857 | +def test_cluster_up_with_write_to_file(mocker): |
| 858 | + """Test cluster.up() method with write_to_file enabled""" |
| 859 | + mocker.patch("kubernetes.config.load_kube_config", return_value="ignore") |
| 860 | + mocker.patch("codeflare_sdk.ray.cluster.cluster.Cluster._throw_for_no_raycluster") |
| 861 | + mocker.patch("codeflare_sdk.ray.cluster.cluster.config_check") |
| 862 | + mocker.patch( |
| 863 | + "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object", |
| 864 | + return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"), |
| 865 | + ) |
| 866 | + |
| 867 | + # Mock the create_namespaced_custom_object |
| 868 | + mock_create = MagicMock() |
| 869 | + mocker.patch( |
| 870 | + "kubernetes.client.CustomObjectsApi.create_namespaced_custom_object", |
| 871 | + return_value=mock_create, |
| 872 | + ) |
| 873 | + |
| 874 | + # Create cluster with write_to_file=True |
| 875 | + config = ClusterConfiguration( |
| 876 | + name="test-cluster-up", |
| 877 | + namespace="ns", |
| 878 | + num_workers=1, |
| 879 | + worker_cpu_requests=1, |
| 880 | + worker_cpu_limits=1, |
| 881 | + worker_memory_requests=2, |
| 882 | + worker_memory_limits=2, |
| 883 | + write_to_file=True, |
| 884 | + appwrapper=True, |
| 885 | + ) |
| 886 | + |
| 887 | + cluster = Cluster(config) |
| 888 | + |
| 889 | + # Mock file reading |
| 890 | + import tempfile |
| 891 | + import os |
| 892 | + |
| 893 | + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: |
| 894 | + f.write("apiVersion: v1\nkind: ConfigMap\nmetadata:\n name: test") |
| 895 | + temp_file = f.name |
| 896 | + |
| 897 | + try: |
| 898 | + cluster.resource_yaml = temp_file |
| 899 | + cluster.up() |
| 900 | + finally: |
| 901 | + os.unlink(temp_file) |
| 902 | + |
| 903 | + |
773 | 904 | # Make sure to always keep this function last |
774 | 905 | def test_cleanup(): |
775 | 906 | os.remove(f"{aw_dir}test-all-params.yaml") |
|
0 commit comments