|
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 |
| 43 | +import ray |
| 44 | +import tempfile |
42 | 45 |
|
43 | 46 | parent = Path(__file__).resolve().parents[4] # project directory |
44 | 47 | expected_clusters_dir = f"{parent}/tests/test_cluster_yamls" |
@@ -377,8 +380,6 @@ def test_cluster_uris(mocker): |
377 | 380 |
|
378 | 381 |
|
379 | 382 | def test_ray_job_wrapping(mocker): |
380 | | - import ray |
381 | | - |
382 | 383 | def ray_addr(self, *args): |
383 | 384 | return self._address |
384 | 385 |
|
@@ -770,6 +771,189 @@ 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_component_resources_with_write_to_file(mocker): |
| 858 | + """Test _component_resources_up with write_to_file enabled""" |
| 859 | + mocker.patch("kubernetes.config.load_kube_config", return_value="ignore") |
| 860 | + mocker.patch( |
| 861 | + "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object", |
| 862 | + return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"), |
| 863 | + ) |
| 864 | + |
| 865 | + # Mock the _create_resources function |
| 866 | + mocker.patch("codeflare_sdk.ray.cluster.cluster._create_resources") |
| 867 | + |
| 868 | + # Create cluster with write_to_file=True (without appwrapper) |
| 869 | + config = ClusterConfiguration( |
| 870 | + name="test-cluster-component", |
| 871 | + namespace="ns", |
| 872 | + num_workers=1, |
| 873 | + worker_cpu_requests=1, |
| 874 | + worker_cpu_limits=1, |
| 875 | + worker_memory_requests=2, |
| 876 | + worker_memory_limits=2, |
| 877 | + write_to_file=True, |
| 878 | + appwrapper=False, |
| 879 | + ) |
| 880 | + |
| 881 | + cluster = Cluster(config) |
| 882 | + |
| 883 | + # Mock file reading and test _component_resources_up |
| 884 | + |
| 885 | + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: |
| 886 | + f.write("apiVersion: v1\nkind: ConfigMap\nmetadata:\n name: test") |
| 887 | + temp_file = f.name |
| 888 | + |
| 889 | + try: |
| 890 | + mock_api = MagicMock() |
| 891 | + cluster.resource_yaml = temp_file |
| 892 | + cluster._component_resources_up("ns", mock_api) |
| 893 | + # If we got here without error, the write_to_file path was executed |
| 894 | + assert True |
| 895 | + finally: |
| 896 | + os.unlink(temp_file) |
| 897 | + |
| 898 | + |
| 899 | +def test_get_cluster_status_functions(mocker): |
| 900 | + """Test _app_wrapper_status and _ray_cluster_status functions""" |
| 901 | + from codeflare_sdk.ray.cluster.cluster import ( |
| 902 | + _app_wrapper_status, |
| 903 | + _ray_cluster_status, |
| 904 | + ) |
| 905 | + |
| 906 | + mocker.patch("kubernetes.config.load_kube_config", return_value="ignore") |
| 907 | + mocker.patch("codeflare_sdk.ray.cluster.cluster.config_check") |
| 908 | + |
| 909 | + # Test _app_wrapper_status when cluster not found |
| 910 | + mocker.patch( |
| 911 | + "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object", |
| 912 | + return_value={"items": []}, |
| 913 | + ) |
| 914 | + result = _app_wrapper_status("non-existent-cluster", "ns") |
| 915 | + assert result is None |
| 916 | + |
| 917 | + # Test _ray_cluster_status when cluster not found |
| 918 | + mocker.patch( |
| 919 | + "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object", |
| 920 | + return_value={"items": []}, |
| 921 | + ) |
| 922 | + result = _ray_cluster_status("non-existent-cluster", "ns") |
| 923 | + assert result is None |
| 924 | + |
| 925 | + |
| 926 | +def test_cluster_namespace_type_error(mocker): |
| 927 | + """Test TypeError when namespace is not a string""" |
| 928 | + mocker.patch("kubernetes.config.load_kube_config", return_value="ignore") |
| 929 | + mocker.patch( |
| 930 | + "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object", |
| 931 | + return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"), |
| 932 | + ) |
| 933 | + |
| 934 | + # Mock get_current_namespace to return a non-string value (e.g., int) |
| 935 | + mocker.patch( |
| 936 | + "codeflare_sdk.ray.cluster.cluster.get_current_namespace", return_value=12345 |
| 937 | + ) |
| 938 | + |
| 939 | + config = ClusterConfiguration( |
| 940 | + name="test-cluster-type-error", |
| 941 | + namespace=None, # Will trigger namespace check |
| 942 | + num_workers=1, |
| 943 | + worker_cpu_requests=1, |
| 944 | + worker_cpu_limits=1, |
| 945 | + worker_memory_requests=2, |
| 946 | + worker_memory_limits=2, |
| 947 | + ) |
| 948 | + |
| 949 | + # This should raise TypeError because get_current_namespace returns int |
| 950 | + with pytest.raises( |
| 951 | + TypeError, |
| 952 | + match="Namespace 12345 is of type.*Check your Kubernetes Authentication", |
| 953 | + ): |
| 954 | + Cluster(config) |
| 955 | + |
| 956 | + |
773 | 957 | # Make sure to always keep this function last |
774 | 958 | def test_cleanup(): |
775 | 959 | os.remove(f"{aw_dir}test-all-params.yaml") |
|
0 commit comments