| 475 | |
| 476 | #[test] |
| 477 | fn test_extract_class() { |
| 478 | let mut parser = PythonParser::new(); |
| 479 | let source = r#" |
| 480 | class UserService: |
| 481 | def __init__(self, db): |
| 482 | self.db = db |
| 483 | |
| 484 | def get_user(self, user_id: str) -> dict: |
| 485 | return self.db.find(user_id) |
| 486 | |
| 487 | def _private_method(self): |
| 488 | pass |
| 489 | "#; |
| 490 | let entities = parser.extract(source, "service.py"); |
| 491 | |
| 492 | let service = entities |
| 493 | .iter() |
| 494 | .find(|e| e.name == "UserService") |
| 495 | .expect("Should find UserService"); |
| 496 | assert_eq!(service.kind, EntityKind::Class); |
| 497 | assert!(service.exported); |
| 498 | |
| 499 | let init = entities.iter().find(|e| e.name == "__init__"); |
| 500 | assert!(init.is_some(), "Should find __init__"); |
| 501 | assert_eq!(init.unwrap().kind, EntityKind::Method); |
| 502 | |
| 503 | let get_user = entities.iter().find(|e| e.name == "get_user"); |
| 504 | assert!(get_user.is_some(), "Should find get_user"); |
| 505 | assert_eq!(get_user.unwrap().kind, EntityKind::Method); |
| 506 | |
| 507 | let private = entities.iter().find(|e| e.name == "_private_method"); |
| 508 | assert!(private.is_some(), "Should find _private_method"); |
| 509 | assert!( |
| 510 | !private.unwrap().exported, |
| 511 | "Private method should not be exported" |
| 512 | ); |
| 513 | } |
| 514 | |
| 515 | #[test] |
| 516 | fn test_extract_class_with_bases() { |