added more tests and some refactors

marirs · Dec 4, 2022 · 2f3a66e · 2f3a66e
1 parent 9fe554e
commit 2f3a66e
Show file tree

Hide file tree

Showing 6 changed files with 129 additions and 14 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "ip2location"
-version = "0.3.2"
+version = "0.3.3"
 authors = ["MARIRS <[email protected]>"]
 description = "Find geo information & proxy information based on the given IP using IP2Location BIN databases"
 keywords = ["ip2location", "geoip", "geolocation", "ip", "proxy"]

diff --git a/README.md b/README.md
@@ -35,6 +35,64 @@ cargo t -v
 ip2location = "0.3.2"
 ```
 
+### Example
+```rust
+use crate::{error, Record, DB};
+
+const IPV4BIN: &str = "data/IP2LOCATION-LITE-DB1.BIN";
+const IPV6BIN: &str = "data/IP2LOCATION-LITE-DB1.IPV6.BIN";
+const IP2PROXYBIN: &str = "data/IP2PROXY-IP-COUNTRY.BIN";
+
+// Lookup an IP v4 in the IP2Location V6 BIN Database
+fn ip_lookup_in_ipv6bin() -> Result<(), error::Error> {
+    let mut db = DB::from_file(IPV6BIN)?;
+    let record = db.ip_lookup("43.224.159.155".parse().unwrap())?;
+    let record = if let Record::LocationDb(rec) = record {
+        Some(rec)
+    } else {
+        None
+    };
+    assert!(record.is_some());
+    let record = record.unwrap();
+    assert!(!record.country.is_none());
+    assert_eq!(record.country.clone().unwrap().short_name, "IN");
+    assert_eq!(record.country.unwrap().long_name, "India");
+    Ok(())
+}
+
+// Lookup an IP v4 in the IP2Location V4 BIN Database
+fn ip_lookup_in_ipv4bin() -> Result<(), error::Error> {
+    let mut db = DB::from_file(IPV4BIN)?;
+    let record = db.ip_lookup("43.224.159.155".parse().unwrap())?;
+    let record = if let Record::LocationDb(rec) = record {
+        Some(rec)
+    } else {
+        None
+    };
+    assert!(record.is_some());
+    let record = record.unwrap();
+    assert!(!record.country.is_none());
+    assert_eq!(record.country.clone().unwrap().short_name, "IN");
+    assert_eq!(record.country.unwrap().long_name, "India");
+    Ok(())
+}
+
+// Lookup an IP in the Proxy Database
+fn ip_lookup_in_proxy_bin() -> Result<(), error::Error> {
+    let mut db = DB::from_file(IP2PROXYBIN)?;
+    let record = db.ip_lookup("1.1.1.1".parse().unwrap())?;
+    let record = if let Record::ProxyDb(rec) = record {
+        Some(rec)
+    } else {
+        None
+    };
+    assert!(record.is_some());
+    let record = record.unwrap();
+    assert!(!record.country.is_none());
+    Ok(())
+}
+```
+
 ### Executing the Example
 ```bash
 cargo b --example

diff --git a/src/common.rs b/src/common.rs
@@ -175,6 +175,12 @@ impl DB {
         //!
         //! let mut db = DB::from_file_mmap("data/IP2PROXY-IP-COUNTRY.BIN").unwrap();
         //!```
+        if !path.as_ref().exists() {
+            return Err(Error::IoError(
+                "Error opening DB file: No such file or directory".to_string(),
+            ));
+        }
+
         if let Ok(location_db) = LocationDB::from_file_mmap(&path) {
             Ok(DB::LocationDb(location_db))
         } else if let Ok(proxy_db) = ProxyDB::from_file_mmap(&path) {

diff --git a/src/ip2location/db.rs b/src/ip2location/db.rs
@@ -66,10 +66,16 @@ impl LocationDB {
         //!
         //! let mut db = LocationDB::from_file("data/IP2LOCATION-LITE-DB1.BIN").unwrap();
         //!```
+        if !path.as_ref().exists() {
+            return Err(Error::IoError(
+                "Error opening DB file: No such file or directory".to_string(),
+            ));
+        }
+
         let db = File::open(&path)?;
-        let mut ss = Self::new(Source::File(path.as_ref().to_path_buf(), db));
-        ss.read_header()?;
-        Ok(ss)
+        let mut ldb = Self::new(Source::File(path.as_ref().to_path_buf(), db));
+        ldb.read_header()?;
+        Ok(ldb)
     }
 
     pub fn from_file_mmap<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
@@ -83,11 +89,17 @@ impl LocationDB {
         //!
         //! let mut db = DB::from_file_mmap("data/IP2LOCATION-LITE-DB1.BIN").unwrap();
         //!```
+        if !path.as_ref().exists() {
+            return Err(Error::IoError(
+                "Error opening DB file: No such file or directory".to_string(),
+            ));
+        }
+
         let db = File::open(&path)?;
         let mm = unsafe { Mmap::map(&db) }?;
-        let mut ss = Self::new(Source::Mmap(path.as_ref().to_path_buf(), mm));
-        ss.read_header()?;
-        Ok(ss)
+        let mut ldb = Self::new(Source::Mmap(path.as_ref().to_path_buf(), mm));
+        ldb.read_header()?;
+        Ok(ldb)
     }
 
     pub fn print_db_info(&self) {

diff --git a/src/ip2proxy/db.rs b/src/ip2proxy/db.rs
@@ -65,10 +65,16 @@ impl ProxyDB {
         //!
         //! let mut db = DB::from_file("data/IP2PROXY-IP-COUNTRY.BIN").unwrap();
         //!```
+        if !path.as_ref().exists() {
+            return Err(Error::IoError(
+                "Error opening DB file: No such file or directory".to_string(),
+            ));
+        }
+
         let db = File::open(&path)?;
-        let mut ss = Self::new(Source::File(path.as_ref().to_path_buf(), db));
-        ss.read_header()?;
-        Ok(ss)
+        let mut pdb = Self::new(Source::File(path.as_ref().to_path_buf(), db));
+        pdb.read_header()?;
+        Ok(pdb)
     }
 
     pub fn from_file_mmap<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
@@ -82,11 +88,17 @@ impl ProxyDB {
         //!
         //! let mut db = DB::from_file_mmap("data/IP2PROXY-IP-COUNTRY.BIN").unwrap();
         //!```
+        if !path.as_ref().exists() {
+            return Err(Error::IoError(
+                "Error opening DB file: No such file or directory".to_string(),
+            ));
+        }
+
         let db = File::open(&path)?;
         let mm = unsafe { Mmap::map(&db) }?;
-        let mut ss = Self::new(Source::Mmap(path.as_ref().to_path_buf(), mm));
-        ss.read_header()?;
-        Ok(ss)
+        let mut pdb = Self::new(Source::Mmap(path.as_ref().to_path_buf(), mm));
+        pdb.read_header()?;
+        Ok(pdb)
     }
 
     pub fn ip_lookup(&mut self, ip: IpAddr) -> Result<ProxyRecord, Error> {

diff --git a/src/tests/tests_lib.rs b/src/tests/tests_lib.rs
@@ -2,6 +2,7 @@ use crate::{error, Record, DB};
 
 const IPV4BIN: &str = "data/IP2LOCATION-LITE-DB1.BIN";
 const IPV6BIN: &str = "data/IP2LOCATION-LITE-DB1.IPV6.BIN";
+const IP2PROXYBIN: &str = "data/IP2PROXY-IP-COUNTRY.BIN";
 
 #[test]
 fn test_ipv4_lookup_in_ipv4bin() -> Result<(), error::Error> {
@@ -106,7 +107,7 @@ fn test_ipv6_lookup_using_mmap() -> Result<(), error::Error> {
 }
 
 #[test]
-fn test_err_filenotfound() -> Result<(), error::Error> {
+fn test_err_filenotfound_location() -> Result<(), error::Error> {
     let db = DB::from_file("nonexistant.bin");
     assert!(db.is_err());
     let result = &db.unwrap_err();
@@ -115,3 +116,29 @@ fn test_err_filenotfound() -> Result<(), error::Error> {
     assert_eq!(result, expected);
     Ok(())
 }
+
+#[test]
+fn test_err_filenotfound_mmap_location() -> Result<(), error::Error> {
+    let db = DB::from_file_mmap("nonexistant.bin");
+    assert!(db.is_err());
+    let result = &db.unwrap_err();
+    let expected =
+        &error::Error::IoError("Error opening DB file: No such file or directory".to_string());
+    assert_eq!(result, expected);
+    Ok(())
+}
+
+#[test]
+fn test_ip_lookup_in_proxy_bin() -> Result<(), error::Error> {
+    let mut db = DB::from_file(IP2PROXYBIN)?;
+    let record = db.ip_lookup("1.1.1.1".parse().unwrap())?;
+    let record = if let Record::ProxyDb(rec) = record {
+        Some(rec)
+    } else {
+        None
+    };
+    assert!(record.is_some());
+    let record = record.unwrap();
+    assert!(!record.country.is_none());
+    Ok(())
+}