edgefirst_client/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright © 2025 Au-Zone Technologies. All Rights Reserved.
3
4// SPDX-License-Identifier: Apache-2.0
5// Copyright © 2025 Au-Zone Technologies. All Rights Reserved.
6
7//! # EdgeFirst Studio Client Library
8//!
9//! The EdgeFirst Studio Client Library provides a Rust client for interacting
10//! with EdgeFirst Studio, a comprehensive platform for computer vision and
11//! machine learning workflows. This library enables developers to
12//! programmatically manage datasets, annotations, training sessions, and other
13//! Studio resources.
14//!
15//! ## Features
16//!
17//! - **Authentication**: Secure token-based authentication with automatic
18//!   renewal
19//! - **Dataset Management**: Upload, download, and manage datasets with various
20//!   file types
21//! - **Annotation Management**: Create, update, and retrieve annotations for
22//!   computer vision tasks
23//! - **Training & Validation**: Manage machine learning training and validation
24//!   sessions
25//! - **Project Organization**: Organize work into projects with hierarchical
26//!   structure
27//! - **Polars Integration**: Optional integration with Polars DataFrames for
28//!   data analysis
29//!
30//! ## Quick Start
31//!
32//! ```rust,no_run
33//! use edgefirst_client::{Client, Error};
34//!
35//! #[tokio::main]
36//! async fn main() -> Result<(), Error> {
37//!     // Create a new client
38//!     let client = Client::new()?;
39//!
40//!     // Authenticate with username and password
41//!     let client = client.with_login("username", "password").await?;
42//!
43//!     // List available projects
44//!     let projects = client.projects(None).await?;
45//!     println!("Found {} projects", projects.len());
46//!
47//!     Ok(())
48//! }
49//! ```
50//!
51//! ## Optional Features
52//!
53//! - `polars`: Enables integration with Polars DataFrames for enhanced data
54//!   manipulation
55
56mod api;
57mod client;
58mod dataset;
59mod error;
60pub mod format;
61mod retry;
62mod storage;
63
64pub use crate::{
65    api::{
66        AnnotationSetID, AppId, Artifact, DatasetID, DatasetParams, Experiment, ExperimentID,
67        ImageId, Organization, OrganizationID, Parameter, PresignedUrl, Project, ProjectID,
68        SampleID, SamplesCountResult, SamplesPopulateParams, SamplesPopulateResult, SequenceId,
69        Snapshot, SnapshotFromDatasetResult, SnapshotID, SnapshotRestoreResult, Stage, Task,
70        TaskID, TaskInfo, TrainingSession, TrainingSessionID, ValidationSession,
71        ValidationSessionID,
72    },
73    client::{Client, Progress},
74    dataset::{
75        Annotation, AnnotationSet, AnnotationType, Box2d, Box3d, Dataset, FileType, GpsData,
76        ImuData, Label, Location, Mask, Sample, SampleFile,
77    },
78    error::Error,
79    retry::{RetryScope, classify_url},
80    storage::{FileTokenStorage, MemoryTokenStorage, StorageError, TokenStorage},
81};
82
83#[cfg(feature = "polars")]
84#[allow(deprecated)] // Re-exported for backwards compatibility
85pub use crate::dataset::annotations_dataframe;
86
87#[cfg(feature = "polars")]
88pub use crate::dataset::samples_dataframe;
89
90#[cfg(feature = "polars")]
91pub use crate::dataset::unflatten_polygon_coordinates;
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use std::{
97        collections::HashMap,
98        env,
99        fs::{File, read_to_string},
100        io::Write,
101        path::PathBuf,
102    };
103
104    /// Get the test data directory (target/testdata)
105    /// Creates it if it doesn't exist
106    fn get_test_data_dir() -> PathBuf {
107        let test_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
108            .parent()
109            .expect("CARGO_MANIFEST_DIR should have parent")
110            .parent()
111            .expect("workspace root should exist")
112            .join("target")
113            .join("testdata");
114
115        std::fs::create_dir_all(&test_dir).expect("Failed to create test data directory");
116        test_dir
117    }
118
119    #[ctor::ctor]
120    fn init() {
121        env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
122    }
123
124    async fn get_client() -> Result<Client, Error> {
125        let client = Client::new()?.with_token_path(None)?;
126
127        let client = match env::var("STUDIO_TOKEN") {
128            Ok(token) => client.with_token(&token)?,
129            Err(_) => client,
130        };
131
132        let client = match env::var("STUDIO_SERVER") {
133            Ok(server) => client.with_server(&server)?,
134            Err(_) => client,
135        };
136
137        let client = match (env::var("STUDIO_USERNAME"), env::var("STUDIO_PASSWORD")) {
138            (Ok(username), Ok(password)) => client.with_login(&username, &password).await?,
139            _ => client,
140        };
141
142        client.verify_token().await?;
143
144        Ok(client)
145    }
146
147    /// Helper: Get training session for "Unit Testing" project
148    async fn get_training_session_for_artifacts() -> Result<TrainingSession, Error> {
149        let client = get_client().await?;
150        let project = client
151            .projects(Some("Unit Testing"))
152            .await?
153            .into_iter()
154            .next()
155            .ok_or_else(|| Error::InvalidParameters("Unit Testing project not found".into()))?;
156        let experiment = client
157            .experiments(project.id(), Some("Unit Testing"))
158            .await?
159            .into_iter()
160            .next()
161            .ok_or_else(|| Error::InvalidParameters("Unit Testing experiment not found".into()))?;
162        let session = client
163            .training_sessions(experiment.id(), Some("modelpack-960x540"))
164            .await?
165            .into_iter()
166            .next()
167            .ok_or_else(|| {
168                Error::InvalidParameters("modelpack-960x540 session not found".into())
169            })?;
170        Ok(session)
171    }
172
173    /// Helper: Get training session for "modelpack-usermanaged"
174    async fn get_training_session_for_checkpoints() -> Result<TrainingSession, Error> {
175        let client = get_client().await?;
176        let project = client
177            .projects(Some("Unit Testing"))
178            .await?
179            .into_iter()
180            .next()
181            .ok_or_else(|| Error::InvalidParameters("Unit Testing project not found".into()))?;
182        let experiment = client
183            .experiments(project.id(), Some("Unit Testing"))
184            .await?
185            .into_iter()
186            .next()
187            .ok_or_else(|| Error::InvalidParameters("Unit Testing experiment not found".into()))?;
188        let session = client
189            .training_sessions(experiment.id(), Some("modelpack-usermanaged"))
190            .await?
191            .into_iter()
192            .next()
193            .ok_or_else(|| {
194                Error::InvalidParameters("modelpack-usermanaged session not found".into())
195            })?;
196        Ok(session)
197    }
198
199    #[tokio::test]
200    async fn test_training_session() -> Result<(), Error> {
201        let client = get_client().await?;
202        let project = client.projects(Some("Unit Testing")).await?;
203        assert!(!project.is_empty());
204        let project = project
205            .first()
206            .expect("'Unit Testing' project should exist");
207        let experiment = client
208            .experiments(project.id(), Some("Unit Testing"))
209            .await?;
210        let experiment = experiment
211            .first()
212            .expect("'Unit Testing' experiment should exist");
213
214        let sessions = client
215            .training_sessions(experiment.id(), Some("modelpack-usermanaged"))
216            .await?;
217        assert_ne!(sessions.len(), 0);
218        let session = sessions
219            .first()
220            .expect("Training sessions should exist for experiment");
221
222        let metrics = HashMap::from([
223            ("epochs".to_string(), Parameter::Integer(10)),
224            ("loss".to_string(), Parameter::Real(0.05)),
225            (
226                "model".to_string(),
227                Parameter::String("modelpack".to_string()),
228            ),
229        ]);
230
231        session.set_metrics(&client, metrics).await?;
232        let updated_metrics = session.metrics(&client).await?;
233        assert_eq!(updated_metrics.len(), 3);
234        assert_eq!(updated_metrics.get("epochs"), Some(&Parameter::Integer(10)));
235        assert_eq!(updated_metrics.get("loss"), Some(&Parameter::Real(0.05)));
236        assert_eq!(
237            updated_metrics.get("model"),
238            Some(&Parameter::String("modelpack".to_string()))
239        );
240
241        println!("Updated Metrics: {:?}", updated_metrics);
242
243        let mut labels = tempfile::NamedTempFile::new()?;
244        write!(labels, "background")?;
245        labels.flush()?;
246
247        session
248            .upload(
249                &client,
250                &[(
251                    "artifacts/labels.txt".to_string(),
252                    labels.path().to_path_buf(),
253                )],
254            )
255            .await?;
256
257        let labels = session.download(&client, "artifacts/labels.txt").await?;
258        assert_eq!(labels, "background");
259
260        Ok(())
261    }
262
263    #[tokio::test]
264    async fn test_validate() -> Result<(), Error> {
265        let client = get_client().await?;
266        let project = client.projects(Some("Unit Testing")).await?;
267        assert!(!project.is_empty());
268        let project = project
269            .first()
270            .expect("'Unit Testing' project should exist");
271
272        let sessions = client.validation_sessions(project.id()).await?;
273        for session in &sessions {
274            let s = client.validation_session(session.id()).await?;
275            assert_eq!(s.id(), session.id());
276            assert_eq!(s.description(), session.description());
277        }
278
279        let session = sessions
280            .into_iter()
281            .find(|s| s.name() == "modelpack-usermanaged")
282            .ok_or_else(|| {
283                Error::InvalidParameters(format!(
284                    "Validation session 'modelpack-usermanaged' not found in project '{}'",
285                    project.name()
286                ))
287            })?;
288
289        let metrics = HashMap::from([("accuracy".to_string(), Parameter::Real(0.95))]);
290        session.set_metrics(&client, metrics).await?;
291
292        let metrics = session.metrics(&client).await?;
293        assert_eq!(metrics.get("accuracy"), Some(&Parameter::Real(0.95)));
294
295        Ok(())
296    }
297
298    #[tokio::test]
299    async fn test_download_artifact_success() -> Result<(), Error> {
300        let trainer = get_training_session_for_artifacts().await?;
301        let client = get_client().await?;
302        let artifacts = client.artifacts(trainer.id()).await?;
303        assert!(!artifacts.is_empty());
304
305        let test_dir = get_test_data_dir();
306        let artifact = &artifacts[0];
307        let output_path = test_dir.join(artifact.name());
308
309        client
310            .download_artifact(
311                trainer.id(),
312                artifact.name(),
313                Some(output_path.clone()),
314                None,
315            )
316            .await?;
317
318        assert!(output_path.exists());
319        if output_path.exists() {
320            std::fs::remove_file(&output_path)?;
321        }
322
323        Ok(())
324    }
325
326    #[tokio::test]
327    async fn test_download_artifact_not_found() -> Result<(), Error> {
328        let trainer = get_training_session_for_artifacts().await?;
329        let client = get_client().await?;
330        let test_dir = get_test_data_dir();
331        let fake_path = test_dir.join("nonexistent_artifact.txt");
332
333        let result = client
334            .download_artifact(
335                trainer.id(),
336                "nonexistent_artifact.txt",
337                Some(fake_path.clone()),
338                None,
339            )
340            .await;
341
342        assert!(result.is_err());
343        assert!(!fake_path.exists());
344
345        Ok(())
346    }
347
348    #[tokio::test]
349    async fn test_artifacts() -> Result<(), Error> {
350        let client = get_client().await?;
351        let project = client.projects(Some("Unit Testing")).await?;
352        assert!(!project.is_empty());
353        let project = project
354            .first()
355            .expect("'Unit Testing' project should exist");
356        let experiment = client
357            .experiments(project.id(), Some("Unit Testing"))
358            .await?;
359        let experiment = experiment
360            .first()
361            .expect("'Unit Testing' experiment should exist");
362        let trainer = client
363            .training_sessions(experiment.id(), Some("modelpack-960x540"))
364            .await?;
365        let trainer = trainer
366            .first()
367            .expect("'modelpack-960x540' training session should exist");
368        let artifacts = client.artifacts(trainer.id()).await?;
369        assert!(!artifacts.is_empty());
370
371        let test_dir = get_test_data_dir();
372
373        for artifact in artifacts {
374            let output_path = test_dir.join(artifact.name());
375            client
376                .download_artifact(
377                    trainer.id(),
378                    artifact.name(),
379                    Some(output_path.clone()),
380                    None,
381                )
382                .await?;
383
384            // Clean up downloaded file
385            if output_path.exists() {
386                std::fs::remove_file(&output_path)?;
387            }
388        }
389
390        let fake_path = test_dir.join("fakefile.txt");
391        let res = client
392            .download_artifact(trainer.id(), "fakefile.txt", Some(fake_path.clone()), None)
393            .await;
394        assert!(res.is_err());
395        assert!(!fake_path.exists());
396
397        Ok(())
398    }
399
400    #[tokio::test]
401    async fn test_download_checkpoint_success() -> Result<(), Error> {
402        let trainer = get_training_session_for_checkpoints().await?;
403        let client = get_client().await?;
404        let test_dir = get_test_data_dir();
405
406        // Create temporary test file
407        let checkpoint_path = test_dir.join("test_checkpoint.txt");
408        {
409            let mut f = File::create(&checkpoint_path)?;
410            f.write_all(b"Test Checkpoint Content")?;
411        }
412
413        // Upload the checkpoint
414        trainer
415            .upload(
416                &client,
417                &[(
418                    "checkpoints/test_checkpoint.txt".to_string(),
419                    checkpoint_path.clone(),
420                )],
421            )
422            .await?;
423
424        // Download and verify
425        let download_path = test_dir.join("downloaded_checkpoint.txt");
426        client
427            .download_checkpoint(
428                trainer.id(),
429                "test_checkpoint.txt",
430                Some(download_path.clone()),
431                None,
432            )
433            .await?;
434
435        let content = read_to_string(&download_path)?;
436        assert_eq!(content, "Test Checkpoint Content");
437
438        // Cleanup
439        if checkpoint_path.exists() {
440            std::fs::remove_file(&checkpoint_path)?;
441        }
442        if download_path.exists() {
443            std::fs::remove_file(&download_path)?;
444        }
445
446        Ok(())
447    }
448
449    #[tokio::test]
450    async fn test_download_checkpoint_not_found() -> Result<(), Error> {
451        let trainer = get_training_session_for_checkpoints().await?;
452        let client = get_client().await?;
453        let test_dir = get_test_data_dir();
454        let fake_path = test_dir.join("nonexistent_checkpoint.txt");
455
456        let result = client
457            .download_checkpoint(
458                trainer.id(),
459                "nonexistent_checkpoint.txt",
460                Some(fake_path.clone()),
461                None,
462            )
463            .await;
464
465        assert!(result.is_err());
466        assert!(!fake_path.exists());
467
468        Ok(())
469    }
470
471    #[tokio::test]
472    async fn test_checkpoints() -> Result<(), Error> {
473        let client = get_client().await?;
474        let project = client.projects(Some("Unit Testing")).await?;
475        assert!(!project.is_empty());
476        let project = project
477            .first()
478            .expect("'Unit Testing' project should exist");
479        let experiment = client
480            .experiments(project.id(), Some("Unit Testing"))
481            .await?;
482        let experiment = experiment.first().ok_or_else(|| {
483            Error::InvalidParameters(format!(
484                "Experiment 'Unit Testing' not found in project '{}'",
485                project.name()
486            ))
487        })?;
488        let trainer = client
489            .training_sessions(experiment.id(), Some("modelpack-usermanaged"))
490            .await?;
491        let trainer = trainer
492            .first()
493            .expect("'modelpack-usermanaged' training session should exist");
494
495        let test_dir = get_test_data_dir();
496        let checkpoint_path = test_dir.join("checkpoint.txt");
497        let checkpoint2_path = test_dir.join("checkpoint2.txt");
498
499        {
500            let mut chkpt = File::create(&checkpoint_path)?;
501            chkpt.write_all(b"Test Checkpoint")?;
502        }
503
504        trainer
505            .upload(
506                &client,
507                &[(
508                    "checkpoints/checkpoint.txt".to_string(),
509                    checkpoint_path.clone(),
510                )],
511            )
512            .await?;
513
514        client
515            .download_checkpoint(
516                trainer.id(),
517                "checkpoint.txt",
518                Some(checkpoint2_path.clone()),
519                None,
520            )
521            .await?;
522
523        let chkpt = read_to_string(&checkpoint2_path)?;
524        assert_eq!(chkpt, "Test Checkpoint");
525
526        let fake_path = test_dir.join("fakefile.txt");
527        let res = client
528            .download_checkpoint(trainer.id(), "fakefile.txt", Some(fake_path.clone()), None)
529            .await;
530        assert!(res.is_err());
531        assert!(!fake_path.exists());
532
533        // Clean up
534        if checkpoint_path.exists() {
535            std::fs::remove_file(&checkpoint_path)?;
536        }
537        if checkpoint2_path.exists() {
538            std::fs::remove_file(&checkpoint2_path)?;
539        }
540
541        Ok(())
542    }
543
544    #[tokio::test]
545    async fn test_task_retrieval() -> Result<(), Error> {
546        let client = get_client().await?;
547
548        // Test: Get all tasks
549        let tasks = client.tasks(None, None, None, None).await?;
550        assert!(!tasks.is_empty());
551
552        // Test: Get task info for first task
553        let task_id = tasks[0].id();
554        let task_info = client.task_info(task_id).await?;
555        assert_eq!(task_info.id(), task_id);
556
557        Ok(())
558    }
559
560    #[tokio::test]
561    async fn test_task_filtering_by_name() -> Result<(), Error> {
562        let client = get_client().await?;
563        let project = client.projects(Some("Unit Testing")).await?;
564        let project = project
565            .first()
566            .expect("'Unit Testing' project should exist");
567
568        // Test: Get tasks by name
569        let tasks = client
570            .tasks(Some("modelpack-usermanaged"), None, None, None)
571            .await?;
572
573        if !tasks.is_empty() {
574            // Get detailed info for each task
575            let task_infos = tasks
576                .into_iter()
577                .map(|t| client.task_info(t.id()))
578                .collect::<Vec<_>>();
579            let task_infos = futures::future::try_join_all(task_infos).await?;
580
581            // Filter by project
582            let filtered = task_infos
583                .into_iter()
584                .filter(|t| t.project_id() == Some(project.id()))
585                .collect::<Vec<_>>();
586
587            if !filtered.is_empty() {
588                assert_eq!(filtered[0].project_id(), Some(project.id()));
589            }
590        }
591
592        Ok(())
593    }
594
595    #[tokio::test]
596    async fn test_task_status_and_stages() -> Result<(), Error> {
597        let client = get_client().await?;
598
599        // Get first available task
600        let tasks = client.tasks(None, None, None, None).await?;
601        if tasks.is_empty() {
602            return Ok(());
603        }
604
605        let task_id = tasks[0].id();
606
607        // Test: Get task status
608        let status = client.task_status(task_id, "training").await?;
609        assert_eq!(status.id(), task_id);
610        assert_eq!(status.status(), "training");
611
612        // Test: Set stages
613        let stages = [
614            ("download", "Downloading Dataset"),
615            ("train", "Training Model"),
616            ("export", "Exporting Model"),
617        ];
618        client.set_stages(task_id, &stages).await?;
619
620        // Test: Update stage
621        client
622            .update_stage(task_id, "download", "running", "Downloading dataset", 50)
623            .await?;
624
625        // Verify task with updated stages
626        let updated_task = client.task_info(task_id).await?;
627        assert_eq!(updated_task.id(), task_id);
628
629        Ok(())
630    }
631
632    #[tokio::test]
633    async fn test_tasks() -> Result<(), Error> {
634        let client = get_client().await?;
635        let tasks = client.tasks(None, None, None, None).await?;
636
637        for task in tasks {
638            let task_info = client.task_info(task.id()).await?;
639            println!("{} - {}", task, task_info);
640        }
641
642        let tasks = client
643            .tasks(Some("modelpack-usermanaged"), None, None, None)
644            .await?;
645        let tasks = tasks
646            .into_iter()
647            .map(|t| client.task_info(t.id()))
648            .collect::<Vec<_>>();
649        let tasks = futures::future::try_join_all(tasks).await?;
650        assert_ne!(tasks.len(), 0);
651        let task = &tasks[0];
652
653        let t = client.task_status(task.id(), "training").await?;
654        assert_eq!(t.id(), task.id());
655        assert_eq!(t.status(), "training");
656
657        let stages = [
658            ("download", "Downloading Dataset"),
659            ("train", "Training Model"),
660            ("export", "Exporting Model"),
661        ];
662        client.set_stages(task.id(), &stages).await?;
663
664        client
665            .update_stage(task.id(), "download", "running", "Downloading dataset", 50)
666            .await?;
667
668        let task = client.task_info(task.id()).await?;
669        println!("task progress: {:?}", task.stages());
670
671        Ok(())
672    }
673
674    // ============================================================================
675    // Retry URL Classification Tests
676    // ============================================================================
677
678    mod retry_url_classification {
679        use super::*;
680
681        #[test]
682        fn test_studio_api_base_url() {
683            // Base production URL
684            assert_eq!(
685                classify_url("https://edgefirst.studio/api"),
686                RetryScope::StudioApi
687            );
688        }
689
690        #[test]
691        fn test_studio_api_with_trailing_slash() {
692            // Trailing slash should be handled correctly
693            assert_eq!(
694                classify_url("https://edgefirst.studio/api/"),
695                RetryScope::StudioApi
696            );
697        }
698
699        #[test]
700        fn test_studio_api_with_path() {
701            // API endpoints with additional path segments
702            assert_eq!(
703                classify_url("https://edgefirst.studio/api/datasets"),
704                RetryScope::StudioApi
705            );
706            assert_eq!(
707                classify_url("https://edgefirst.studio/api/auth.login"),
708                RetryScope::StudioApi
709            );
710            assert_eq!(
711                classify_url("https://edgefirst.studio/api/trainer/session"),
712                RetryScope::StudioApi
713            );
714        }
715
716        #[test]
717        fn test_studio_api_with_query_params() {
718            // Query parameters should not affect classification
719            assert_eq!(
720                classify_url("https://edgefirst.studio/api?foo=bar"),
721                RetryScope::StudioApi
722            );
723            assert_eq!(
724                classify_url("https://edgefirst.studio/api/datasets?page=1&limit=10"),
725                RetryScope::StudioApi
726            );
727        }
728
729        #[test]
730        fn test_studio_api_subdomains() {
731            // Server-specific instances (test, stage, saas, ocean, etc.)
732            assert_eq!(
733                classify_url("https://test.edgefirst.studio/api"),
734                RetryScope::StudioApi
735            );
736            assert_eq!(
737                classify_url("https://stage.edgefirst.studio/api"),
738                RetryScope::StudioApi
739            );
740            assert_eq!(
741                classify_url("https://saas.edgefirst.studio/api"),
742                RetryScope::StudioApi
743            );
744            assert_eq!(
745                classify_url("https://ocean.edgefirst.studio/api"),
746                RetryScope::StudioApi
747            );
748        }
749
750        #[test]
751        fn test_studio_api_with_standard_port() {
752            // Standard HTTPS port (443) should be handled
753            assert_eq!(
754                classify_url("https://edgefirst.studio:443/api"),
755                RetryScope::StudioApi
756            );
757            assert_eq!(
758                classify_url("https://test.edgefirst.studio:443/api"),
759                RetryScope::StudioApi
760            );
761        }
762
763        #[test]
764        fn test_studio_api_with_custom_port() {
765            // Custom ports should be handled correctly
766            assert_eq!(
767                classify_url("https://test.edgefirst.studio:8080/api"),
768                RetryScope::StudioApi
769            );
770            assert_eq!(
771                classify_url("https://edgefirst.studio:8443/api"),
772                RetryScope::StudioApi
773            );
774        }
775
776        #[test]
777        fn test_studio_api_http_protocol() {
778            // HTTP (not HTTPS) should still be recognized
779            assert_eq!(
780                classify_url("http://edgefirst.studio/api"),
781                RetryScope::StudioApi
782            );
783            assert_eq!(
784                classify_url("http://test.edgefirst.studio/api"),
785                RetryScope::StudioApi
786            );
787        }
788
789        #[test]
790        fn test_file_io_s3_urls() {
791            // S3 URLs for file operations
792            assert_eq!(
793                classify_url("https://s3.amazonaws.com/bucket/file.bin"),
794                RetryScope::FileIO
795            );
796            assert_eq!(
797                classify_url("https://s3.us-west-2.amazonaws.com/mybucket/data.zip"),
798                RetryScope::FileIO
799            );
800        }
801
802        #[test]
803        fn test_file_io_cloudfront_urls() {
804            // CloudFront URLs for file distribution
805            assert_eq!(
806                classify_url("https://d123abc.cloudfront.net/file.bin"),
807                RetryScope::FileIO
808            );
809            assert_eq!(
810                classify_url("https://d456def.cloudfront.net/path/to/file.tar.gz"),
811                RetryScope::FileIO
812            );
813        }
814
815        #[test]
816        fn test_file_io_non_api_studio_paths() {
817            // Non-API paths on edgefirst.studio domain
818            assert_eq!(
819                classify_url("https://edgefirst.studio/docs"),
820                RetryScope::FileIO
821            );
822            assert_eq!(
823                classify_url("https://edgefirst.studio/download_model"),
824                RetryScope::FileIO
825            );
826            assert_eq!(
827                classify_url("https://test.edgefirst.studio/download_model"),
828                RetryScope::FileIO
829            );
830            assert_eq!(
831                classify_url("https://stage.edgefirst.studio/download_checkpoint"),
832                RetryScope::FileIO
833            );
834        }
835
836        #[test]
837        fn test_file_io_generic_urls() {
838            // Generic download URLs
839            assert_eq!(
840                classify_url("https://example.com/download"),
841                RetryScope::FileIO
842            );
843            assert_eq!(
844                classify_url("https://cdn.example.com/files/data.json"),
845                RetryScope::FileIO
846            );
847        }
848
849        #[test]
850        fn test_security_malicious_url_substring() {
851            // Security: URL with edgefirst.studio in path should NOT match
852            assert_eq!(
853                classify_url("https://evil.com/test.edgefirst.studio/api"),
854                RetryScope::FileIO
855            );
856            assert_eq!(
857                classify_url("https://attacker.com/edgefirst.studio/api/fake"),
858                RetryScope::FileIO
859            );
860        }
861
862        #[test]
863        fn test_edge_case_similar_domains() {
864            // Similar but different domains should be FileIO
865            assert_eq!(
866                classify_url("https://edgefirst.studio.com/api"),
867                RetryScope::FileIO
868            );
869            assert_eq!(
870                classify_url("https://notedgefirst.studio/api"),
871                RetryScope::FileIO
872            );
873            assert_eq!(
874                classify_url("https://edgefirststudio.com/api"),
875                RetryScope::FileIO
876            );
877        }
878
879        #[test]
880        fn test_edge_case_invalid_urls() {
881            // Invalid URLs should default to FileIO
882            assert_eq!(classify_url("not a url"), RetryScope::FileIO);
883            assert_eq!(classify_url(""), RetryScope::FileIO);
884            assert_eq!(
885                classify_url("ftp://edgefirst.studio/api"),
886                RetryScope::FileIO
887            );
888        }
889
890        #[test]
891        fn test_edge_case_url_normalization() {
892            // URL normalization edge cases
893            assert_eq!(
894                classify_url("https://EDGEFIRST.STUDIO/api"),
895                RetryScope::StudioApi
896            );
897            assert_eq!(
898                classify_url("https://test.EDGEFIRST.studio/api"),
899                RetryScope::StudioApi
900            );
901        }
902
903        #[test]
904        fn test_comprehensive_subdomain_coverage() {
905            // Ensure all known server instances are recognized
906            let subdomains = vec![
907                "test", "stage", "saas", "ocean", "prod", "dev", "qa", "demo",
908            ];
909
910            for subdomain in subdomains {
911                let url = format!("https://{}.edgefirst.studio/api", subdomain);
912                assert_eq!(
913                    classify_url(&url),
914                    RetryScope::StudioApi,
915                    "Failed for subdomain: {}",
916                    subdomain
917                );
918            }
919        }
920
921        #[test]
922        fn test_api_path_variations() {
923            // Various API path patterns
924            assert_eq!(
925                classify_url("https://edgefirst.studio/api"),
926                RetryScope::StudioApi
927            );
928            assert_eq!(
929                classify_url("https://edgefirst.studio/api/"),
930                RetryScope::StudioApi
931            );
932            assert_eq!(
933                classify_url("https://edgefirst.studio/api/v1"),
934                RetryScope::StudioApi
935            );
936            assert_eq!(
937                classify_url("https://edgefirst.studio/api/v2/datasets"),
938                RetryScope::StudioApi
939            );
940
941            // Non-/api paths should be FileIO
942            assert_eq!(
943                classify_url("https://edgefirst.studio/apis"),
944                RetryScope::FileIO
945            );
946            assert_eq!(
947                classify_url("https://edgefirst.studio/v1/api"),
948                RetryScope::FileIO
949            );
950        }
951
952        #[test]
953        fn test_port_range_coverage() {
954            // Test various port numbers
955            let ports = vec![80, 443, 8080, 8443, 3000, 5000, 9000];
956
957            for port in ports {
958                let url = format!("https://test.edgefirst.studio:{}/api", port);
959                assert_eq!(
960                    classify_url(&url),
961                    RetryScope::StudioApi,
962                    "Failed for port: {}",
963                    port
964                );
965            }
966        }
967
968        #[test]
969        fn test_complex_query_strings() {
970            // Complex query parameters with special characters
971            assert_eq!(
972                classify_url("https://edgefirst.studio/api?token=abc123&redirect=/dashboard"),
973                RetryScope::StudioApi
974            );
975            assert_eq!(
976                classify_url("https://test.edgefirst.studio/api?q=search%20term&page=1"),
977                RetryScope::StudioApi
978            );
979        }
980
981        #[test]
982        fn test_url_with_fragment() {
983            // URLs with fragments (#) - fragments are not sent to server
984            assert_eq!(
985                classify_url("https://edgefirst.studio/api#section"),
986                RetryScope::StudioApi
987            );
988            assert_eq!(
989                classify_url("https://test.edgefirst.studio/api/datasets#results"),
990                RetryScope::StudioApi
991            );
992        }
993    }
994}