Small dataset added

JuliaWeb · Mar 23, 2021 · 363c955 · 363c955
1 parent a918785
commit 363c955
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 26 deletions.
diff --git a/src/data.jl b/src/data.jl
@@ -77,20 +77,20 @@ function update()
     global dataloaded = false
 end
 
-function load()
-    blockfile = joinpath(DATADIR, BLOCKCSVGZ)
-    locfile = joinpath(DATADIR, CITYCSVGZ)
+function load(datadir = DATADIR)
+    blockfile = joinpath(datadir, BLOCKCSVGZ)
+    locfile = joinpath(datadir, CITYCSVGZ)
 
     local blocks
     local locs
     try
         blocks = GZip.open(blockfile, "r") do stream
-            CSV.File(stream) |> DataFrame
+            CSV.File(read(stream)) |> DataFrame
             # CSV.File(stream, types=[String, Int, Int, String, Int, Int, String, Float64, Float64, Int]) |> DataFrame
         end
         locs = GZip.open(locfile, "r") do stream
             # CSV.File(stream, types=[Int, String, String, String, String, String, String, String, String, String, String, Int, String, Int]) |> DataFrame
-            CSV.File(stream) |> DataFrame
+            CSV.File(read(stream)) |> DataFrame
         end
     catch
         @error "Geolocation data cannot be read. Data directory may be corrupt..."

diff --git a/src/geoip-module.jl b/src/geoip-module.jl
@@ -53,7 +53,7 @@ function geolocate(ip::IPv4; noupdate = true)
     end
 
     # TODO: sentinel value should be returned
-    retdict = Dict{Symbol, Any}()
+    retdict = Dict{String, Any}()
     if (found > 0) && ip in geodata[found, :v4net]
         # Placeholder, should be removed
         row = geodata[found, :]
@@ -62,13 +62,5 @@ function geolocate(ip::IPv4; noupdate = true)
     return retdict
 end
 
-function geolocate(iparr::AbstractArray; noupdate = true)
-    masterdict = Dict{Symbol, Any}[]
-    for el in iparr
-        push!(masterdict, geolocate(el; noupdate = noupdate))
-    end
-    return masterdict
-end
-
 geolocate(ipstr::AbstractString; noupdate = true) = geolocate(IPv4(ipstr); noupdate = noupdate)
 geolocate(ipint::Integer; noupdate = true) = geolocate(IPv4(ipint); noupdate = noupdate)
diff --git a/test/data/GeoLite2-City-Blocks-IPv4.csv.gz b/test/data/GeoLite2-City-Blocks-IPv4.csv.gz
diff --git a/test/data/GeoLite2-City-Locations-en.csv.gz b/test/data/GeoLite2-City-Locations-en.csv.gz
diff --git a/test/test01_base.jl b/test/test01_base.jl
@@ -1,7 +1,7 @@
 module TestBase
 
 using GeoIP
-using Sockets: IPv4
+using Sockets: IPv4, @ip_str
 using Test
 
 # WARNING!!
@@ -10,24 +10,27 @@ using Test
 # Placeholders are left, so they can be used later for actual tests
 # after resolving https://github.com/JuliaWeb/GeoIP.jl/issues/12
 
+TEST_DATADIR = joinpath(dirname(@__FILE__), "data")
+load(TEST_DATADIR)
+
 @testset "Known result" begin
-    ip1 = IPv4("1.2.3.4")
-    # geoip1 = geolocate(ip1; noupdate=false)
-    # @test geoip1[:country_iso_code] == "US"
-    # @test geoip1[:metro_code] == 819
-    # @test ceil(Int, geoip1[:location].x) == -122
+    ip1 = IPv4("1.0.8.1")
+    geoip1 = geolocate(ip1)
+    @test geoip1["country_iso_code"] == "CN"
+    @test geoip1["time_zone"] == "Asia/Shanghai"
+    @test ceil(Int, geoip1["location"].x) == 114
 end
 
 @testset "Null results" begin
-    # @test isempty(geolocate(ip"0.0.0.0"))
-    # @test isempty(geolocate(ip"127.0.0.1"))
+    @test isempty(geolocate(ip"0.0.0.0"))
+    @test isempty(geolocate(ip"127.0.0.1"))
 end
 
 @testset "Array of ip's" begin
-    # result = geolocate([ip"1.2.3.4", ip"8.8.8.8"])
-    # @test length(Set(result)) == 2
-    # @test !isempty(result[1])
-    # @test !isempty(result[2])  
+    result = geolocate.([ip"1.0.16.1", ip"1.0.8.1"])
+    @test length(Set(result)) == 2
+    @test !isempty(result[1])
+    @test !isempty(result[2])  
 end
 
 end # module